/* (C) 2003 XDA Developers  itsme@xs4all.nl
 *
 * $Header$
 */
#include <winbase.h>
#include "debug.h"

#include <map>
#include <vector>
#include <string>
#include <algorithm>
using namespace std;

int g_verbose= 0;

//-------------------------------------------------------------------------------
//  memory map tracker
class MemoryBlock {
public:
    DWORD length;
    DWORD vstart;
    DWORD vend;
    DWORD pstart;
    DWORD pend;
	BYTE *data;

    /* no destructor, to make memory management without copy constructors and refcounting easy */
    /*    I currently don't care about the data buffer leaking */

    bool InPhysRange(DWORD poffset) const { return (pstart <= poffset && poffset < pend); }
    bool InVirtRange(DWORD voffset) const { return (vstart <= voffset && voffset < vend); }

    bool operator <(const MemoryBlock& mb) { return pstart < mb.pstart; }
    bool operator <(DWORD poffset) { return pend < poffset; }
};

typedef vector<MemoryBlock> MemoryBlockVector;

class MemoryMap {
public:
    bool MapMemory(DWORD voffset, DWORD poffset, DWORD size);
    
    void *PhysToVirt(DWORD poffset) const;
    DWORD GetOfs(void *ptr) const;
    BYTE GetPhysByte(DWORD poffset) const;
    DWORD GetPhysDword(DWORD poffset) const;

    DWORD FirstPhysAddress() const;
    DWORD LastPhysAddress() const;
    DWORD NextValidAddress(DWORD address) const;

    void SerializeBlocks();

	void DumpPhysToVirtMap();
	void DumpVirtToPhysMap();
private:
    MemoryBlockVector m_blocks;
};
MemoryMap  g_vmem;	// maps virtual -> physical addresses


class MemRegion {
public:
    DWORD start;
    DWORD end;
    DWORD length;
    string *description;

    /* no destructor, to make memory management without copy constructors and refcounting easy */
    /*    I currently don't care about the description buffer leaking */
    MemRegion(DWORD start, DWORD end) : start(start), end(end), description(NULL), length(end-start) {}

    bool operator <(const MemRegion& r) const { return start < r.start || (start==r.start && length<r.length); }

};

typedef vector<MemRegion> MemRegionVector;
class MemRegions {
public:
    MemRegion& MarkRange(DWORD start, DWORD end, const char *msg, ...);
    MemRegion& MarkRegion(DWORD start, DWORD length, const char *msg, ...);
    MemRegion& MarkRegion_v(DWORD start, DWORD length, const char *msg, va_list ap);

    void DumpMemoryMap();

private:
    MemRegionVector m_list;

};

//--------------------------- global variables
MemRegions g_vregions;

// maps physical memory back to virtual
bool MemoryMap::MapMemory(DWORD voffset, DWORD poffset, DWORD size)
{
    MemoryBlock mb;

    mb.length= size;
    mb.vstart= voffset;
    mb.vend= voffset+size;
    mb.pstart= poffset;
    mb.pend= poffset+size;

	mb.data= (BYTE*)mb.vstart;

    // keep m_blocks sorted.
    MemoryBlockVector::iterator i;
    for (i=m_blocks.begin() ; i!=m_blocks.end(); ++i)
        if (mb.pstart < (*i).pstart)
            break;

    m_blocks.insert(i, mb);

    return true;
}	
void MemoryMap::SerializeBlocks()
{
    for (MemoryBlockVector::iterator i= m_blocks.begin(); i!=m_blocks.end() ; ++i)
    {
        for (MemoryBlockVector::iterator j= i+1 ;
                j!=m_blocks.end()
                && (*i).pend == (*j).pstart
                && (*i).vend == (*j).vstart ;
                j= i+1)
        {
            if ((*i).pstart >= (*j).pstart)
                debug("WARNING: m_blocks not sorted as expected\n");

            (*i).pend    = (*j).pend;
            (*i).vend    = (*j).vend;
            (*i).length += (*j).length;
            m_blocks.erase(j);

        }
    }
}
void MemoryMap::DumpPhysToVirtMap()
{
    for (MemoryBlockVector::iterator i= m_blocks.begin(); i!=m_blocks.end() ; ++i)
    {
    }
}

void MemoryMap::DumpVirtToPhysMap()
{
}

BYTE MemoryMap::GetPhysByte(DWORD poffset) const
{
    BYTE *p= (BYTE*)PhysToVirt(poffset);
    if (p==NULL)
        return 0;
    return *p;
}
DWORD MemoryMap::GetPhysDword(DWORD poffset) const
{
    DWORD *p= (DWORD*)PhysToVirt(poffset);
    if (p==NULL)
        return 0;
    return *p;
}
void *MemoryMap::PhysToVirt(DWORD poffset) const
{
    for (MemoryBlockVector::const_iterator i=m_blocks.begin() ; i!=m_blocks.end(); ++i)
    {
        if ((*i).InPhysRange(poffset)) {
            return (*i).data+(poffset-(*i).pstart);
        }
    }
    debug("ERROR: could not find pointer for phys ofs %08lx\n", poffset);
    return NULL;
}
DWORD MemoryMap::GetOfs(void *ptr) const
{
    for (MemoryBlockVector::const_iterator i=m_blocks.begin() ; i!= m_blocks.end() ; ++i)
    {
        if ((*i).data <= ptr && ptr < (*i).data+(*i).length)
        {
            return ((BYTE*)ptr - (*i).data) + (*i).pstart;
        }
    }
    debug("ERROR: could not find offset for ptr %08lx\n", ptr);
    return 0;
}

DWORD MemoryMap::FirstPhysAddress() const
{
    MemoryBlockVector::const_iterator i=m_blocks.begin();

    return (*i).pstart;
}
DWORD MemoryMap::LastPhysAddress() const
{
    MemoryBlockVector::const_reverse_iterator i=m_blocks.rbegin();
    return (*i).pend;
}

//=========================================================

MemRegion& MemRegions::MarkRange(DWORD start, DWORD end, const char *msg, ...)
{
    va_list ap;
    va_start(ap, msg);
    MemRegion& r= MarkRegion_v(start, end-start, msg, ap);
    va_end(ap);

    return r;
}
MemRegion& MemRegions::MarkRegion(DWORD start, DWORD length, const char *msg, ...)
{
    va_list ap;
    va_start(ap, msg);
    MemRegion& r= MarkRegion_v(start, length, msg, ap);
    va_end(ap);

    return r;
}
MemRegion& MemRegions::MarkRegion_v(DWORD start, DWORD length, const char *msg, va_list ap)
{
    char msgbuf[1024];
    _vsnprintf(msgbuf, 1024, msg, ap);

    MemRegion *m= new MemRegion(start, start+length);
    m->description= new string(msgbuf);
    if (m->description==NULL)
    {
        debug("error allocating memory\n");
    }

    m_list.push_back(*m);

    return *m;
}

void MemRegions::DumpMemoryMap()
{
	debug("marked %d memory entries\n", m_list.size());
    sort(m_list.begin(), m_list.end());

    DWORD offset= 0;
    for (MemRegionVector::iterator i=m_list.begin() ; i!=m_list.end() ; ++i)
    {
        if (offset < (*i).start) {
            MemRegion m(offset, (*i).start);
            if ( ((*i).start & 3)==0 && (*i).start - offset < 4)
            {
                if (g_verbose>0) {
                    debug("\t%08lx - %08lx L%08lx alignment", m.start, m.end, m.length);
                    //if (m.FirstNonzero()!=m.end)
                    //    bytedump(m.start, m.end);
                }
            }
            else
            {
/*
                DWORD firstnz= max(m.start, m.FirstNonzero() & ~3);
                DWORD lastnz= min(m.end, (m.LastNonzero() & ~3)+4);
                if (firstnz==m.end)
                    debug("\n%08lx - %08lx L%08lx NUL", m.start, m.end, m.length);
                else {
                    if (firstnz != m.start)
                        debug("\n%08lx - %08lx L%08lx NUL", m.start, firstnz, firstnz-m.start);
*/
                    //debug("\n%08lx - %08lx L%08lx unknown", firstnz, lastnz, lastnz-firstnz);
				debug("\n%08lx - %08lx L%08lx unknown", m.start, m.end, m.end-m.start);
/*
                    if (lastnz-firstnz<16)
                        bytedump(firstnz, lastnz);
                    else if (lastnz-firstnz<64)
                        dworddump(firstnz, lastnz);
                    if (lastnz != m.end)
                        debug("\n%08lx - %08lx L%08lx NUL", lastnz, m.end, m.end-lastnz);
                }
*/
            }
        }
        else if (offset > (*i).start) {
            debug("\n!!! overlap of %08lx bytes\n", offset-(*i).start );
        }

        debug("\n%08lx - %08lx L%08lx %s", (*i).start, (*i).end, (*i).length, (*i).description->c_str());

        offset= (*i).end;
    }

    if (offset<=0xffffffff)
    {
        debug("\n%08lx - %08lx unknown", offset, 0);
    }
    debug("\n");
}

//-------------------------------------------------------------------------------
#define GETBITS(w, start, end) ((w>>end)&((1<<(start-end+1)) - 1))


void DumpTLB2(DWORD sectiontype, DWORD section, DWORD sectionvbase, DWORD basepaddress, DWORD nentries);
void DumpIgnoreTLB2Entry(DWORD sectiontype, DWORD section, DWORD page, DWORD *ppageentry);
void DumpTinyPageDescriptor(DWORD sectiontype, DWORD section, DWORD page, DWORD *ppageentry);
void DumpSmallPageDescriptor(DWORD sectiontype, DWORD section, DWORD page, DWORD *ppageentry);
void DumpLargePageDescriptor(DWORD sectiontype, DWORD section, DWORD page, DWORD *ppageentry);

void DumpTLB(DWORD TLBBase);
void DumpCoarseTLBEntry(DWORD section, DWORD sectionvbase, DWORD *psectionentry);
void DumpFineTLBEntry(DWORD section, DWORD sectionvbase, DWORD *psectionentry);
void DumpIgnoreTLBEntry(DWORD section, DWORD sectionvbase, DWORD *psectionentry);
void DumpSectionDescriptor(DWORD section, DWORD sectionvbase, DWORD *psectionentry);

extern "C" BOOL SetKMode(BOOL bFlag);
extern "C" DWORD SetProcPermissions(DWORD dwPerms);


int WINAPI WinMain(	HINSTANCE hInstance,
					HINSTANCE hPrevInstance,
					LPTSTR    lpCmdLine,
					int       nCmdShow)
{
    DebugSetLogfile("tlbdump.log");

    g_vmem.MapMemory(0x80000000, 0x00000000, 0x02000000);
    g_vmem.MapMemory(0x83000000, 0x10000000, 0x00100000);
    g_vmem.MapMemory(0x8c000000, 0x10200000, 0x01c00000);
    g_vmem.MapMemory(0x8f100000, 0x11d00000, 0x00300000);

    g_vmem.MapMemory(0xfffd0000, 0x10200000, 0x00004000);

    BOOL bMode = SetKMode(TRUE);
    DWORD dwPerm = SetProcPermissions(0xFFFFFFFF);

    DumpTLB(_tcstoul(lpCmdLine, 0, 0));

    SetKMode(bMode );
    SetProcPermissions(dwPerm);


	debug("virtual memory map-----------------------------------\n");
    g_vregions.DumpMemoryMap();

	debug("reverse memory map-----------------------------------\n");
	g_vmem.DumpPhysToVirtMap();

	return 0;
}


void DumpTLB(DWORD tlbbase)
{
	// this is from p15, c2

	// set it with 'MCR p15, 0, Rn, c2, c0'
	// read it with 'MRC p15, 0, Rn, c2, c0'
	//  or EE?2nf10   ?=0: MCR, ?=1: MRC

    DWORD *TLB1= (DWORD*)g_vmem.PhysToVirt(tlbbase&0xffffc000);
    if (TLB1==NULL) {
        debug("DumpTLB: invalid tlbbase\n");
        return;
    }

	g_vregions.MarkRegion((DWORD)TLB1, 0, "start of top level pagetable");
	g_vregions.MarkRegion((DWORD)TLB1+4096*sizeof(DWORD), 0, "end of top level pagetable");
    for (int section=0 ; section<4096 ; section++)
    {
        DWORD sectionentry= TLB1[section];
        DWORD sectionvbase= section<<20;

        switch(sectionentry&3)
        {
        case 0: DumpIgnoreTLBEntry(section, sectionvbase, TLB1+section); break;
        case 1: DumpCoarseTLBEntry(section, sectionvbase, TLB1+section); break;
        case 2: DumpSectionDescriptor(section, sectionvbase, TLB1+section); break;
        case 3: DumpFineTLBEntry(section, sectionvbase, TLB1+section); break;
        }
    }
}

void DumpIgnoreTLBEntry(DWORD section, DWORD sectionvbase, DWORD *psectionentry)
{
    //debug("DumpIgnoreTLBEntry(%08lx %08lx %08lx)\n", section, sectionvbase, psectionentry);
    if (*psectionentry&~3)
		g_vregions.MarkRegion((DWORD)psectionentry, sizeof(DWORD), "IGN %08lx  : %08lx", sectionvbase, *psectionentry);
}

// describes 1MB page

// physaddr= bits( TLB1[bits(addr,31,20)], 31, 20)<<20 | bits(addr,19,0)
void DumpSectionDescriptor(DWORD section, DWORD sectionvbase, DWORD *psectionentry)
{
    //debug("DumpSectionDescriptor(%08lx %08lx %08lx)\n", section, sectionvbase, psectionentry);
	DWORD sectionentry = *psectionentry;

    DWORD basepaddress= GETBITS(sectionentry, 31, 20)<<20;
    int unused1= GETBITS(sectionentry, 19, 12);
    int accessperms= GETBITS(sectionentry, 11, 10);
    int unused2= GETBITS(sectionentry, 9, 9);
    int domain= GETBITS(sectionentry, 8, 5);
    int imp= GETBITS(sectionentry, 4, 4);
    int cachable_bit= GETBITS(sectionentry, 3, 3);
    int bufferable_bit= GETBITS(sectionentry, 2, 2);

    if (unused1 || unused2)
        debug("unexpected bits not null: %02x %02x\n", unused1, unused2);
	g_vregions.MarkRegion((DWORD)psectionentry, sizeof(DWORD), "1st level tlb: section entry v%08lx -> p%08lx P=%d dom=%x imp=%d cachable=%d bufferable=%d",
		sectionvbase, basepaddress, accessperms, domain, imp, cachable_bit, bufferable_bit);

	g_vregions.MarkRegion(sectionvbase, 1024*1024, "section phys=%08lx", basepaddress);

	g_vmem.MapMemory(sectionvbase, basepaddress, 1024*1024);
}

// TLB2=bits(TLB1[bits(addr,31,20)], 31, 10)<<10 
//    TLB3= TLB2[bits(addr, 19,12)]
//    physaddr= bits(TLB3, ?,?) | bits(addr(11,0)

// can point to large and small pages.
// a coarse table is 1Kbyte
void DumpCoarseTLBEntry(DWORD section, DWORD sectionvbase, DWORD *psectionentry)
{
    //debug("DumpCoarseTLBEntry(%08lx %08lx %08lx)\n", section, sectionvbase, psectionentry);
	DWORD sectionentry = *psectionentry;
    DWORD basepaddress= GETBITS(sectionentry, 31, 10)<<10;
    int unused1= GETBITS(sectionentry, 9, 9);
    int domain= GETBITS(sectionentry, 8, 5);
    int imp= GETBITS(sectionentry, 4, 2);

    if (unused1)
        debug("unexpected bits not null: %02x\n", unused1);

	g_vregions.MarkRegion((DWORD)psectionentry, sizeof(DWORD), "1st level tlb: coarse tlb entry  @%08lx  dom=%x imp=%d   v=%08lx", 
		basepaddress, domain, imp, sectionvbase);
    DumpTLB2(1, section, sectionvbase, basepaddress, 256);
}

// 
// TLB2=bits(TLB1[bits(addr,31,20)], 31, 12)<<12
//    TLB3= TLB2[bits(addr, 19,10)]
//    physaddr= bits(TLB3, ?,?) | bits(addr(9,0)

// a coarse table is 4Kbyte
void DumpFineTLBEntry(DWORD section, DWORD sectionvbase, DWORD *psectionentry)
{
    //debug("DumpFineTLBEntry(%08lx %08lx %08lx)\n", section, sectionvbase, psectionentry);
	DWORD sectionentry = *psectionentry;
    DWORD basepaddress= GETBITS(sectionentry, 31, 12)<<12;
    int unused1= GETBITS(sectionentry, 11, 9);
    int domain= GETBITS(sectionentry, 8, 5);
    int imp= GETBITS(sectionentry, 4, 2);

    if (unused1)
        debug("unexpected bits not null: %02x\n", unused1);

	g_vregions.MarkRegion((DWORD)psectionentry, sizeof(DWORD), "1st level tlb: fine tlb entry @%08lx dom=%x imp=%d  v=%08lx", 
			basepaddress, domain, imp, sectionvbase);

    DumpTLB2(3, section, sectionvbase, basepaddress, 1024);
}

// sectionbase is the 1MB virtual memory range this tlb maps.
// baseaddress is the physical memmory address of this 2nd level tlb
void DumpTLB2(DWORD sectiontype, DWORD section, DWORD sectionvbase, DWORD basepaddress, DWORD nentries)
{
    //debug("DumpTLB2(%08lx %08lx %08lx)\n", sectiontype, section, sectionvbase, basepaddress, nentries);
    DWORD *TLB2= (DWORD*)g_vmem.PhysToVirt(basepaddress);
	if (TLB2==NULL)
	{
		debug("DumpTLB2: invalid basepaddress\n");
		return;
	}
	g_vregions.MarkRegion((DWORD)TLB2, 0, "start 2nd level page table for section %08lx  (v=%08lx)", section, sectionvbase);
	g_vregions.MarkRegion((DWORD)TLB2+nentries*sizeof(DWORD), 0, "end 2nd level page table for section %08lx  (v=%08lx)", section, sectionvbase);

    for (DWORD page=0 ; page<nentries ; page++)
    {
        DWORD pageentry= TLB2[page];

        switch(pageentry&3)
        {
        case 0: DumpIgnoreTLB2Entry(sectiontype, section, page, TLB2+page); break;
        case 1: DumpLargePageDescriptor(sectiontype, section, page, TLB2+page); break;
        case 2: DumpSmallPageDescriptor(sectiontype, section, page, TLB2+page); break;
        case 3: DumpTinyPageDescriptor(sectiontype, section, page, TLB2+page); break;
        }

    }
}

void DumpIgnoreTLB2Entry(DWORD sectiontype, DWORD section, DWORD page, DWORD *ppageentry)
{
    //debug("DumpIgnoreTLB2Entry(%08lx %08lx %08lx)\n", sectiontype, section, page, ppageentry);
    if (*ppageentry&~3)
		g_vregions.MarkRegion((DWORD)ppageentry, sizeof(DWORD), "IGN sect %08lx  page %08lx : %08lx", section, page, *ppageentry);
}

// describes 64K page

// bits(tlb, 31, 16) | bits(addr, 15, 0)
//    
// large page table entry must repeat 16 times in a coarse page table,
//                            and 64 times in a fine page table
// or in other words:
//      in a 2nd level coarse page table, bits 15-12 of the page index are ignored
//  and in a 2nd level fine page table, bits  15-10 of the page index are ignored
void DumpLargePageDescriptor(DWORD sectiontype, DWORD section, DWORD page, DWORD *ppageentry)
{
    //debug("DumpLargePageDescriptor(%08lx %08lx %08lx)\n", sectiontype, section, page, ppageentry);
	DWORD pageentry= *ppageentry;

    DWORD basepaddress= GETBITS(pageentry, 31, 16)<<16;
    DWORD unused1= GETBITS(pageentry, 15, 12);
    DWORD ap3= GETBITS(pageentry, 11, 10);
    DWORD ap2= GETBITS(pageentry, 9, 8);
    DWORD ap1= GETBITS(pageentry, 7, 6);
    DWORD ap0= GETBITS(pageentry, 5, 4);
    int cachable_bit= GETBITS(pageentry, 3, 3);
    int bufferable_bit= GETBITS(pageentry, 2, 2);

    if (unused1)
        debug("unexpected non-0: %02x\n", unused1);

	DWORD pagevaddress= (section<<20) | (page<< (sectiontype==1?12:10))&0xffff0000;

	g_vregions.MarkRegion((DWORD)ppageentry, sizeof(DWORD), 
			"2nd level tlb: large page entry  v=%08lx p=%08lx ap=%d.%d.%d.%d cachable=%d bufferable=%d   [ s=%08lx ]",
			pagevaddress, basepaddress, ap0, ap1, ap2, ap3, cachable_bit, bufferable_bit, section<<20);

	g_vregions.MarkRegion(pagevaddress, 65536, "large page phys=%08lx  entry=%08lx", basepaddress, ppageentry);

	g_vmem.MapMemory(pagevaddress, basepaddress, 65536);
}

// describe 4K page
// bits(tlb, 31, 12) | bits(addr, 11, 0)
void DumpSmallPageDescriptor(DWORD sectiontype, DWORD section, DWORD page, DWORD *ppageentry)
{
    //debug("DumpSmallPageDescriptor(%08lx %08lx %08lx)\n", sectiontype, section, page, ppageentry);
	DWORD pageentry= *ppageentry;

    DWORD basepaddress= GETBITS(pageentry, 31, 12)<<12;
    DWORD ap3= GETBITS(pageentry, 11, 10);
    DWORD ap2= GETBITS(pageentry, 9, 8);
    DWORD ap1= GETBITS(pageentry, 7, 6);
    DWORD ap0= GETBITS(pageentry, 5, 4);
    int cachable_bit= GETBITS(pageentry, 3, 3);
    int bufferable_bit= GETBITS(pageentry, 2, 2);

	DWORD pagevaddress= (section<<20) | (page<<12);

	g_vregions.MarkRegion((DWORD)ppageentry, sizeof(DWORD), 
			"2nd level tlb: small page entry  v=%08lx p=%08lx ap=%d.%d.%d.%d cachable=%d bufferable=%d   [ s=%08lx ]",
			pagevaddress, basepaddress, ap0, ap1, ap2, ap3, cachable_bit, bufferable_bit, section<<20);

	g_vregions.MarkRegion(pagevaddress, 4096, "small page phys=%08lx  entry=%08lx", basepaddress, ppageentry);

	g_vmem.MapMemory(pagevaddress, basepaddress, 4096);

}

// describe a 1K page
// bits(tlb, 31, 10) | bits(addr, 9, 0)
void DumpTinyPageDescriptor(DWORD sectiontype, DWORD section, DWORD page, DWORD *ppageentry)
{
    //debug("DumpTinyPageDescriptor(%08lx %08lx %08lx)\n", sectiontype, section, page, ppageentry);
	DWORD pageentry= *ppageentry;

    DWORD basepaddress= GETBITS(pageentry, 31, 10)<<10;
    DWORD unused1= GETBITS(pageentry, 9, 6);
    DWORD ap= GETBITS(pageentry, 5, 4);
    int cachable_bit= GETBITS(pageentry, 3, 3);
    int bufferable_bit= GETBITS(pageentry, 2, 2);

    if (unused1)
        debug("unexpected non-0: %02x\n", unused1);


	DWORD pagevaddress= (section<<20) | (page<<12);

	g_vregions.MarkRegion((DWORD)ppageentry, sizeof(DWORD), 
			"2nd level tlb: tiny page entry  v=%08lx p=%08lx ap=%d cachable=%d bufferable=%d   [ s=%08lx ]",
			pagevaddress, basepaddress, ap, cachable_bit, bufferable_bit, section<<20);

	g_vregions.MarkRegion(pagevaddress, 1024, "tiny page phys=%08lx  entry=%08lx", basepaddress, ppageentry);

	g_vmem.MapMemory(pagevaddress, basepaddress, 1024);

}