// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.
// ===========================================================================
// File: CEEGEN.H
//
// ===========================================================================

#ifndef _CEEGEN_H_
#define _CEEGEN_H_

#include "cor.h"
#include "iceefilegen.h"
#include "ceegentokenmapper.h"

class CeeSection;
class CeeSectionString;
class CCeeGen;
class PESectionMan;
class PESection;

typedef DWORD StringRef;

#if 0

 This is a description of the current implementation of these types for generating
 CLR modules.
 
  ICeeGen - interface to generate in-memory CLR module.

  CCeeGen - implementation of ICeeGen. Currently it uses both CeeSections
            as well as PESections (inside PESectionMan), and maintains a
            1:1 relation between them. This is ugly.

  CeeFileGenWriter - Provides functionality to write in-memory module to
                     PE format file. Uses PEWriter (m_pSectionMan) for
                     file-writing functionality

  PEWriter - It can generate a PE format file. It also knows to apply
             pointer relocs when it lays out the PESections.


  ICeeFileGen - Interface used by compilers, ngen, etc, to generate
                a CLR file.
                Has a bunch of methods to emit signatures, tokens, methods,
                etc which are not implemented. These are left over from before

  CeeFileGenWriter, PEWriter and ICeeFileGen are implemented in mscorpe.dll.
  The other types live in mscorwks.dll (src\md\CeeFileGen)
  
                                                     +----------------------------+
                                                     |       ICeeGen              |
                                                     |                            |
                                                     |  COM-style version of      |
                                                     |  ICeeFileGen. HCEEFILE is  |
                                                     |  replaced with "this"      |
        +-------------------------+                  |                            |
        |     CeeSectionImpl      |                  +----------------------------+
        +-------------------------+                                |
           |                  |                                    |
           |                  |                                    v
           |                  v                      +---------------------------+
           |             +------------+              |      CCeeGen              |
           |             |            |              +---------------------------+
           |             | CeeSection |  contains    |                           |
           |             |            |<-------------| CeeSection* m_sections    |
           |             +------------+              |                           |
           |                                        /| PESectionMan m_pSectionMan|
           |                                       / |                           |
           |             +-----------------+      /  +---------------------------+
           v             |   PESectionMan  |<----+                 |
 +-----------+           |                 |  contains             |
 | PESection |           +-----------------+                       |
 |           |  contains | PESection *     |                       v
 |           |<----------|      sectStart, |        +------------------------------+
 +-----------+           |      sectCur,   |        |       CeeFileGenWriter       |
                         |      sectEnd    |        +------------------------------+
                         +-----------------+        | Does meta-data specific      |
                                  |                 | stuff and then dispatches to |
                                  |                 | m_pSectionMan.PEWriter::***()|
                                  |                 |                              |
                                  v                 +------------------------------+
                       +------------------------+                  ^
                       |       PEWriter         |                  |wraps
                       +------------------------+                  |
                       | Low -level file writer |    +----------------------------+
                       | Knows how to do        |    |        ICeeFileGen         |
                       | pointer relocs         |    |                            |
                       |                        |    | C-style inteface. Deals    |
                       +------------------------+    | with HCEEFILE, HCEESECTION |
                                                     | etc. It is mostly just a   |
                                                     | thin wrapper for a         |
                                                     | CeeFileGenWriter           |
                                                     +----------------------------+

#endif // 0

// ***** CeeSection classes

class CeeSectionImpl {
  public:
    virtual unsigned dataLen() = 0;
    virtual char * getBlock(
        unsigned len, 
        unsigned align = 1) = 0;
    virtual HRESULT addSectReloc(
        unsigned               offset, 
        CeeSection &           relativeTo, 
        CeeSectionRelocType    reloc = srRelocAbsolute, 
        CeeSectionRelocExtra * extra = NULL) = 0;
    virtual HRESULT addBaseReloc(
        unsigned               offset, 
        CeeSectionRelocType    reloc = srRelocHighLow, 
        CeeSectionRelocExtra * extra = NULL) = 0;
    virtual HRESULT directoryEntry(unsigned num) = 0;
    virtual unsigned char * name() = 0;
    virtual char * computePointer(unsigned offset) const = 0;
    virtual BOOL containsPointer(__in char * ptr) const = 0;
    virtual unsigned computeOffset(__in char * ptr) const = 0;
    virtual unsigned getBaseRVA() = 0;
    virtual void SetInitialGrowth(unsigned growth) = 0;
};

class CeeSection {
    // m_ceeFile allows inter-section communication
    CCeeGen &m_ceeFile;

    // abstract away implementation to allow inheritance from CeeSection
    CeeSectionImpl &m_impl;

  public:
    enum RelocFlags {
        RELOC_NONE = 0,

        // address should be fixed up to be a RVA not a normal address
        RELOC_RVA = 1
    };

    CeeSection(CCeeGen &ceeFile, CeeSectionImpl &impl)
        : m_ceeFile(ceeFile), m_impl(impl) { LIMITED_METHOD_CONTRACT; }

    virtual ~CeeSection() {LIMITED_METHOD_CONTRACT;  }

    // bytes in this section at present
    unsigned dataLen();

    // section base, after linking
    unsigned getBaseRVA();

    // get a block to write on (use instead of write to avoid copy)
    char* getBlock(unsigned len, unsigned align=1);

    // Indicates that the DWORD at 'offset' in the current section should
    // have the base of section 'relativeTo added to it
    HRESULT addSectReloc(unsigned offset, CeeSection& relativeTo,
                         CeeSectionRelocType = srRelocAbsolute, CeeSectionRelocExtra *extra = 0);
    // Add a base reloc for the given offset in the current section
    virtual HRESULT addBaseReloc(unsigned offset, CeeSectionRelocType reloc = srRelocHighLow, CeeSectionRelocExtra *extra = 0);


    // this section will be directory entry 'num'
    HRESULT directoryEntry(unsigned num);

    // return section name
    unsigned char *name();

    // simulate the base + offset with a more complex data storage
    char * computePointer(unsigned offset) const;
    BOOL containsPointer(__in char *ptr) const;
    unsigned computeOffset(__in char *ptr) const;

    CeeSectionImpl &getImpl();
    CCeeGen &ceeFile();
    void SetInitialGrowth(unsigned growth);
};

// ***** CCeeGen class
// Only handles in memory stuff
// Base class for CeeFileGenWriter (which actually generates PEFiles)

class CCeeGen : public ICeeGen, ICeeGenInternal {
    LONG m_cRefs;
    BOOL m_encMode;
  protected:
    short m_textIdx;            // m_sections[] index for the .text section
    short m_metaIdx;            // m_sections[] index for metadata (.text, or .cormeta for obj files)
    short m_corHdrIdx;          // m_sections[] index for the COM+ header (.text0)
    short m_stringIdx;          // m_sections[] index for strings (.text, or .rdata for EnC)
    short m_ilIdx;              // m_sections[] index for IL (.text)
    bool m_objSwitch;

    CeeGenTokenMapper *m_pTokenMap;
    BOOLEAN m_fTokenMapSupported;   // temporary to support both models
    IMapToken *m_pRemapHandler;

    CeeSection **m_sections;
    short m_numSections;
    short m_allocSections;

    PESectionMan * m_peSectionMan;

    IMAGE_COR20_HEADER *m_corHeader;
    DWORD m_corHeaderOffset;

    HRESULT allocateCorHeader();

    HRESULT addSection(CeeSection *section, short *sectionIdx);

    HRESULT setEnCMode();

// Init process: Call static CreateNewInstance() , not operator new
  protected:
    HRESULT Init();
    CCeeGen();

  public:

    virtual ~CCeeGen() {}

    static HRESULT CreateNewInstance(CCeeGen* & pCeeFileGen); // call this to instantiate

    virtual HRESULT Cleanup();

    // ICeeGen interfaces

    ULONG STDMETHODCALLTYPE AddRef();
    ULONG STDMETHODCALLTYPE Release();
    STDMETHODIMP QueryInterface(
        REFIID riid,
        void **ppInterface);

    STDMETHODIMP EmitString (
        __in LPWSTR lpString,               // [IN] String to emit
        ULONG *RVA);

    STDMETHODIMP GetString (
        ULONG RVA,
        __inout LPWSTR *lpString);

    STDMETHODIMP AllocateMethodBuffer (
        ULONG cchBuffer,                    // [IN] Length of string to emit
        UCHAR **lpBuffer,                   // [OUT] Returned buffer
        ULONG *RVA);

    STDMETHODIMP GetMethodBuffer (
        ULONG RVA,
        UCHAR **lpBuffer);

    STDMETHODIMP GetIMapTokenIface (
        IUnknown **pIMapToken);

    STDMETHODIMP GenerateCeeFile ();

    STDMETHODIMP GetIlSection (
        HCEESECTION *section);

    STDMETHODIMP GetStringSection (
        HCEESECTION *section);

    STDMETHODIMP AddSectionReloc (
        HCEESECTION section,
        ULONG offset,
        HCEESECTION relativeTo,
        CeeSectionRelocType relocType);

    STDMETHODIMP GetSectionCreate (
        const char *name,
        DWORD flags,
        HCEESECTION *section);

    STDMETHODIMP GetSectionDataLen (
        HCEESECTION section,
        ULONG *dataLen);

    STDMETHODIMP GetSectionBlock (
        HCEESECTION section,
        ULONG len,
        ULONG align=1,
        void **ppBytes=0);

    STDMETHODIMP TruncateSection (
        HCEESECTION section,
        ULONG len);

    STDMETHODIMP GenerateCeeMemoryImage (void **ppImage);

    STDMETHODIMP ComputePointer (
        HCEESECTION section,
        ULONG RVA,                          // [IN] RVA for method to return
        UCHAR **lpBuffer);                  // [OUT] Returned buffer


    STDMETHODIMP AddNotificationHandler(IUnknown *pHandler);

    // Write the metadata in "emitter" to the default metadata section is "section" is 0
    // If 'section != 0, it will put the data in 'buffer'.  This
    // buffer is assumed to be in 'section' at 'offset' and of size 'buffLen'
    // (should use GetSaveSize to insure that buffer is big enough
    virtual HRESULT emitMetaData(IMetaDataEmit *emitter,
                        CeeSection* section=0, DWORD offset=0, BYTE* buffer=0, unsigned buffLen=0);
    virtual HRESULT getMethodRVA(ULONG codeOffset, ULONG *codeRVA);

    STDMETHODIMP SetInitialGrowth(DWORD growth);

    CeeSection &getTextSection();
    CeeSection &getMetaSection();
    CeeSection &getCorHeaderSection();
    CeeSectionString &getStringSection();
    CeeSection &getIlSection();

    virtual HRESULT getSectionCreate (const char *name, DWORD flags, CeeSection **section=NULL, short *sectionIdx = NULL);

    PESectionMan* getPESectionMan() {
        LIMITED_METHOD_CONTRACT;
        return m_peSectionMan;
    }

    virtual HRESULT getMapTokenIface(IUnknown **pIMapToken, IMetaDataEmit *emitter=0);

    CeeGenTokenMapper *getTokenMapper() {
        LIMITED_METHOD_CONTRACT;
        return m_pTokenMap;
    }

    virtual HRESULT addNotificationHandler(IUnknown *pHandler);

    //Clone is actually a misnomer here.  This method will copy all of the
    //instance variables and then do a deep copy (as necessary) of the sections.
    //Section data will be appended onto any information already in the section.
    //This is done to support the DynamicIL -> PersistedIL transform.
    virtual HRESULT cloneInstance(CCeeGen *destination);

#ifdef EMIT_FIXUPS
public:
    virtual HRESULT addFixup(CeeSection& sectionSource, unsigned offset, CeeSectionRelocType reloc, CeeSection * sectionTarget = NULL, CeeSectionRelocExtra *extra = 0) {
        LIMITED_METHOD_CONTRACT;

        return(E_NOTIMPL);
    }
#endif
};

// ***** CeeSection inline methods

inline unsigned CeeSection::dataLen() {
    WRAPPER_NO_CONTRACT;
    return m_impl.dataLen(); }

inline unsigned CeeSection::getBaseRVA() {
    WRAPPER_NO_CONTRACT;
    return m_impl.getBaseRVA(); }

inline char *CeeSection::getBlock(unsigned len, unsigned align) {
    WRAPPER_NO_CONTRACT;
    return m_impl.getBlock(len, align); }

inline HRESULT CeeSection::addSectReloc(
                unsigned offset, CeeSection& relativeTo, CeeSectionRelocType reloc, CeeSectionRelocExtra *extra) {
    WRAPPER_NO_CONTRACT;
    return(m_impl.addSectReloc(offset, relativeTo, reloc, extra));
}

inline HRESULT CeeSection::addBaseReloc(unsigned offset, CeeSectionRelocType reloc, CeeSectionRelocExtra *extra) {
    WRAPPER_NO_CONTRACT;
    return(m_impl.addBaseReloc(offset, reloc, extra));
}


inline HRESULT CeeSection::directoryEntry(unsigned num) {
    WRAPPER_NO_CONTRACT;
    TESTANDRETURN(num < IMAGE_NUMBEROF_DIRECTORY_ENTRIES, E_INVALIDARG);
    m_impl.directoryEntry(num);
    return S_OK;
}

inline CCeeGen &CeeSection::ceeFile() {
    LIMITED_METHOD_CONTRACT;
    return m_ceeFile; }

inline CeeSectionImpl &CeeSection::getImpl() {
    LIMITED_METHOD_CONTRACT;
    return m_impl; }

inline unsigned char *CeeSection::name() {
    WRAPPER_NO_CONTRACT;
    return m_impl.name();
}

inline char * CeeSection::computePointer(unsigned offset) const
{
    WRAPPER_NO_CONTRACT;
    return m_impl.computePointer(offset);
}

inline BOOL CeeSection::containsPointer(__in char *ptr) const
{
    WRAPPER_NO_CONTRACT;
    return m_impl.containsPointer(ptr);
}

inline unsigned CeeSection::computeOffset(__in char *ptr) const
{
    WRAPPER_NO_CONTRACT;
    return m_impl.computeOffset(ptr);
}

inline void CeeSection::SetInitialGrowth(unsigned growth)
{
    WRAPPER_NO_CONTRACT;
    m_impl.SetInitialGrowth(growth);
}

// ***** CCeeGen inline methods

inline CeeSection &CCeeGen::getTextSection() {
    LIMITED_METHOD_CONTRACT;

    return *m_sections[m_textIdx]; }

inline CeeSection &CCeeGen::getMetaSection() {
    LIMITED_METHOD_CONTRACT;

    return *m_sections[m_metaIdx]; }

inline CeeSection &CCeeGen::getCorHeaderSection() {
    LIMITED_METHOD_CONTRACT;
    _ASSERTE(m_corHdrIdx >= 0);
    return *m_sections[m_corHdrIdx]; }

inline CeeSectionString &CCeeGen::getStringSection() {
    LIMITED_METHOD_CONTRACT;

    return *(CeeSectionString*)m_sections[m_stringIdx]; }

inline CeeSection &CCeeGen::getIlSection() {
    LIMITED_METHOD_CONTRACT;

    return *m_sections[m_ilIdx]; }

#endif