#include <winbase.h>
#include "debug.h"
/* (C) 2003-2007 Willem Jan Hengeveld <itsme@xs4all.nl>
 * Web: http://www.xs4all.nl/~itsme/
 *      http://wiki.xda-developers.com/
 *
 * $Id: dumpmemory.cpp 1502 2007-04-15 07:54:20Z itsme $
 */
#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 start;
    DWORD end;
	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 InRange(DWORD offset) const { return (start <= offset && offset < end); }

    operator <(const MemoryBlock& mb) { return start < mb.start; }
    operator <(DWORD offset) { return end < offset; }
};

typedef vector<MemoryBlock> MemoryBlockVector;

typedef vector<DWORD> DWORDList;
typedef map<DWORD,DWORDList> BlockMapList;

class MemoryMap {
public:
    bool MapMemory(DWORD voffset, DWORD poffset, DWORD size);
    
    void *GetPtr(DWORD offset) const;
    DWORD GetOfs(void *ptr) const;
    BYTE GetByte(DWORD offset) const;
    DWORD GetDword(DWORD offset) const;

    DWORD FirstAddress() const;
    DWORD LastAddress() const;
    DWORD NextValidAddress(DWORD address) const;

	void DumpReverseMap();
	void DumpForwardMap();
private:
    MemoryBlockVector m_blocks;

	BlockMapList m_reversemap;
	BlockMapList m_forwardmap;
};
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) {}

    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.start= poffset;
    mb.end= 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.start < (*i).start)
            break;

    m_blocks.insert(i, mb);

	m_reversemap[poffset].push_back(voffset);
	m_forwardmap[voffset].push_back(poffset);

    return true;
}	
void MemoryMap::DumpReverseMap()
{
	for (BlockMapList::iterator i= m_reversemap.begin() ; i!=m_reversemap.end() ; ++i)
	{
		debug("%08lx :", (*i).first);
		for (DWORDList::iterator j= (*i).second.begin() ; j!=(*i).second.end() ; ++j)
		{
			debug(" %08lx", (*j));
		}
		debug("\n");
	}
}

void MemoryMap::DumpForwardMap()
{
	for (BlockMapList::iterator i= m_forwardmap.begin() ; i!=m_forwardmap.end() ; ++i)
	{
		debug("%08lx :", (*i).first);
		for (DWORDList::iterator j= (*i).second.begin() ; j!=(*i).second.end() ; ++j)
		{
			debug(" %08lx", (*j));
		}
		debug("\n");
	}
}

BYTE MemoryMap::GetByte(DWORD offset) const
{
    BYTE *p= (BYTE*)GetPtr(offset);
    if (p==NULL)
        return 0;
    return *p;
}
DWORD MemoryMap::GetDword(DWORD offset) const
{
    DWORD *p= (DWORD*)GetPtr(offset);
    if (p==NULL)
        return 0;
    return *p;
}
void *MemoryMap::GetPtr(DWORD offset) const
{
    for (MemoryBlockVector::const_iterator i=m_blocks.begin() ; i!=m_blocks.end(); ++i)
    {
        if ((*i).InRange(offset)) {
            return (*i).data+(offset-(*i).start);
        }
    }
    debug("ERROR: could not find pointer for ofs %08lx\n", offset);
    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).start;
        }
    }
    debug("ERROR: could not find offset for ptr %08lx\n", ptr);
    return 0;
}

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

    return (*i).start;
}
DWORD MemoryMap::LastAddress() const
{
    MemoryBlockVector::const_reverse_iterator i=m_blocks.rbegin();
    return (*i).end;
}

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 %d 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)
{
	g_vmem.MapMemory(0xac000000, 0xc0000000, 0x02000000);
	g_vmem.MapMemory(0x8c000000, 0xc0000000, 0x02000000);
	g_vmem.MapMemory(0x80000000, 0x00000000, 0x02000000);

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

	DumpTLB(0xc00b0000);

    SetKMode(bMode );
    SetProcPermissions(dwPerm);


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

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

	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.GetPtr(tlbbase&0xffffc000);

	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)
{
    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)
{
	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)
{
	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)
{
	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)
{
    DWORD *TLB2= (DWORD*)g_vmem.GetPtr(basepaddress);
	if (TLB2==NULL)
	{
		debug(".......\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)
{
    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)
{
	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)
{
	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)
{
	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);

}