#include <vector>
#include <map>
#include <set>
#include <cassert>

#define DEBUG

class InplaceBuilder
{
public:

  InplaceBuilder();
  ~InplaceBuilder();

  void RegisterPointer( void** ptr, unsigned target_size );
  void UnregisterPointer( void** ptr, unsigned target_size );

  template< class T >
  void RegisterPointer( T*& ptr, unsigned target_size )
  {
    RegisterPointer( (void**)&ptr, target_size );
  }

  template< class T >
  void UnregisterPointer( T*& ptr, unsigned target_size )
  {
    UnregisterPointer( (void**)&ptr, target_size );
  }

  template< class T >
  void RegisterRootObject( T* ptr )
  {
    assert( m_root == NULL );
    m_blocks.insert( Block( ptr, sizeof( T ) ) );
    m_root = ptr;
  }

  void Build( std::vector< char >& out ) const;

private:

  struct Block
  {
    Block( void* _s, unsigned _l )
      : start( _s ), length( _l ) {}

    Block()
      : start( NULL ), length( 0 ) {}

    bool operator<( Block const& b ) const
    {
      if( start == b.start )
        return length < b.length;
      else
        return start < b.start;
    }

    void* start;
    unsigned length;
  };

  std::multiset< Block > m_blocks;
  std::set< void** > m_pointers;

  void* m_root;

  void GetFinalBlockList( std::vector< Block >& blocks ) const;
  void* FindContainingBlockBase( void* ptr ) const;

public:
#ifdef DEBUG
  void DebugDump() const;
#endif

};

void FixupInplaceImage( void* buffer, unsigned length );

#include <iostream>

using namespace std;


#define REF_UNSIGNED( p, offset ) (*(unsigned*)(((char*)p) + offset))


InplaceBuilder::InplaceBuilder()
  : m_root( NULL )
{
}

InplaceBuilder::~InplaceBuilder()
{
}

void InplaceBuilder::RegisterPointer( void** ptr, unsigned target_size )
{
  m_pointers.insert( ptr );
  if( target_size != 0 )
    m_blocks.insert( Block( *ptr, target_size ) );
}

void InplaceBuilder::UnregisterPointer( void** ptr, unsigned target_size )
{
  if( target_size != 0 )
  {
    Block b( *ptr, target_size );
    assert( m_blocks.find( b ) != m_blocks.end() );
    m_blocks.erase( m_blocks.find( b ) );
  }
  assert( m_pointers.find( ptr ) != m_pointers.end() );
  m_pointers.erase( ptr );
}

void InplaceBuilder::GetFinalBlockList( vector< Block >& blocks ) const
{
  if( m_root )
    blocks.push_back( Block( m_root, 0 ) );
  for( set< Block >::const_iterator it = m_blocks.begin(); it != m_blocks.end(); ++it )
  {
    bool found = false;
    for( unsigned i = 0; i < blocks.size(); ++i )
    {
      if( it->start == blocks[i].start &&
          it->length > blocks[i].length )
      {
        blocks[i] = *it;
        found = true;
        break;
      }
    }

    if( !found )
      blocks.push_back( *it );
  }
  assert( blocks.size() == 0 || blocks[0].length > 0 );
}

void* InplaceBuilder::FindContainingBlockBase( void* ptr ) const
{
  for( set< Block >::const_iterator it = m_blocks.begin(); it != m_blocks.end(); ++it )
  {
    Block const& b = *it;
    if( ptr >= b.start && (char*)ptr < (char*)b.start + b.length )
      return b.start;
  }
  return NULL;
}

void InplaceBuilder::Build( vector< char >& out ) const
{
  assert( sizeof( unsigned ) == sizeof( void* ) );

  vector< Block > blocks;
  GetFinalBlockList( blocks );  // root guaranteed to be first
  assert( blocks.size() == 0 || blocks[0].start == m_root );

  map< void*, unsigned > block_offsets;

  // determine file offsets for all blocks
  unsigned offset = 0;
  for( unsigned i = 0; i < blocks.size(); ++i )
  {
    block_offsets[blocks[i].start] = offset;
    offset += blocks[i].length;

    // align to 4 bytes
    while( offset & 3 ) ++offset;
  }

  // offset is now the total file length
  out.resize( offset + ( m_pointers.size() * sizeof( unsigned ) ) + sizeof( unsigned ) );

  // add all blocks
  for( unsigned i = 0; i < blocks.size(); ++i )
  {
    memcpy( &out[block_offsets[blocks[i].start]],
            blocks[i].start,
            blocks[i].length );
  }

  // convert pointers
  vector< unsigned > pointer_fixup_table;
  for( set< void** >::const_iterator it = m_pointers.begin(); it != m_pointers.end(); ++it )
  {
    void** pointer_itself = *it;
    void* pointer_target = *pointer_itself;

    // find the block the pointer target is inside
    void* pointer_targets_block = FindContainingBlockBase( pointer_target );
    assert( pointer_targets_block );
    unsigned target_blocks_stored_offset = block_offsets[pointer_targets_block];
    unsigned target_offset_within_block = (char*)pointer_target - (char*)pointer_targets_block;
    unsigned target_value_to_store = target_blocks_stored_offset + target_offset_within_block;

    // find the block the pointer itself is in
    void* pointers_block = FindContainingBlockBase( (void*)pointer_itself );
    assert( pointers_block );
    unsigned pointers_offset_into_block = (char*)pointer_itself - (char*)pointers_block;

    // find the pointer's stored position in the file
    unsigned blocks_stored_offset = block_offsets[pointers_block];
    unsigned pointers_stored_offset = blocks_stored_offset + pointers_offset_into_block;

    // update the stored pointer in place
    *(unsigned*)(&out[pointers_stored_offset]) = target_value_to_store;

    pointer_fixup_table.push_back( pointers_stored_offset );
  }

  // store pointer fixup table and offset
  memcpy( &out[offset], &pointer_fixup_table[0], pointer_fixup_table.size() * sizeof( unsigned ) );
  unsigned table_offset = offset;
  memcpy( &out[out.size()-sizeof( unsigned )], &table_offset, sizeof( unsigned ) );
}


void FixupInplaceImage( void* buffer, unsigned length )
{
  unsigned ptr_table_offset = REF_UNSIGNED( buffer, length - sizeof( unsigned ) );
  unsigned* ptr_table = (unsigned*)(((char*)buffer) + ptr_table_offset);
  unsigned ptr_table_length = (length - ptr_table_offset) / sizeof( unsigned ) - 1;

   for( unsigned i = 0; i < ptr_table_length; ++i )
  {
    unsigned old_ptr = REF_UNSIGNED( buffer, ptr_table[i] );
    unsigned new_ptr = ((unsigned)buffer) + old_ptr;
    REF_UNSIGNED( buffer, ptr_table[i] ) = new_ptr;
    cerr << "DEBUG: fixup " << old_ptr << " -> " << new_ptr << " at " << ptr_table[i] << endl;
  }
}


#ifdef DEBUG
void InplaceBuilder::DebugDump() const
{
  cerr << "Inplace builder debug dump:" << endl;
  cerr << " Blocks:" << endl;
  vector< Block > blocks;
  GetFinalBlockList( blocks );
  for( unsigned i = 0; i < blocks.size(); ++i )
  {
    cerr << "  " << blocks[i].start << " (" << blocks[i].length << " bytes)" << endl;
  }

  cerr << " Pointers:" << endl;
  for( set< void** >::const_iterator it = m_pointers.begin(); it != m_pointers.end(); ++it )
  {
    void** addr = *it;
    void* ptr = *addr;
    Block* b = NULL;
    for( unsigned i = 0; i < blocks.size(); ++i )
      if( blocks[i].start == ptr )
      {
        b = &blocks[i];
        break;
      }
    if( b )
      cerr << "  " << addr << " -> " << b->start << " (" << b->length << " bytes)";
    else
      cerr << "  " << addr << " -> " << "(unknown block)";
    cerr << endl;
  }
}
#endif


struct Test
{
  Test( char* m, Test* n )
  {
    msg = new char[strlen(m)+1];
    strcpy( msg, m );
    next = n;
  }

  char* msg;
  Test* next;

  void RegisterPointers( InplaceBuilder& b )
  {
    b.RegisterPointer( msg, strlen( msg )+1 );
    if( next )
    {
      next->RegisterPointers( b );
      b.RegisterPointer( next, sizeof( Test ) );
    }
  }

  void Display()
  {
    cerr << "Node " << this << ": " << msg << endl;
    if( next )
      next->Display();
  }
};



int main()
{
  Test* t = new Test( "foo", new Test( "bar", new Test( "baz", NULL ) ) );
  InplaceBuilder b;
  t->RegisterPointers( b );
  b.RegisterRootObject( t );
  b.DebugDump();
  vector< char > out;
  b.Build( out );
  cerr << "Built object of " << out.size() << " bytes" << endl;

  FILE* f = fopen( "test.out", "wb" );
  assert( f );
  fwrite( &out[0], out.size(), 1, f );
  fclose( f );

  FixupInplaceImage( &out[0], out.size() );
  Test* fixed_up = (Test*)&out[0];

  fixed_up->Display();

  return 0;
}