compiler/src/cpp_generator.py

import time
import os
import os.path
from string import Template
from parser import *
from generator import *

HEADER_COMMENT = """/**
 * Autogenerated by Thrift
 * ${date}
 *
 * DO NOT EDIT UNLESS YOU ARE SURE THAT YOU KNOW WHAT YOU ARE DOING
 */
 """

CPP_TYPES_HEADER = Template(HEADER_COMMENT+"""
#if !defined(${source}_types_h_)
#define ${source}_types_h_ 1

#include <thrift/Thrift.h>
""")

CPP_TYPES_FOOTER = Template("""
#endif // !defined(${source}_types_h_)
""")

CPP_SERVICES_HEADER = Template(HEADER_COMMENT+"""
#if !defined(${source}_h_)
#define ${source}_h_ 1

#include <thrift/Thrift.h>
""")

CPP_SERVICES_FOOTER = Template("""
#endif // !defined(${source}_h_)""")

def cpp_debug(arg):
    print(arg)

class Indenter(object):
    def __init__(self, level=0, step=4):
	self.level = level
	self.step = step
	self.chunk = ""
	for i in range(step):
	    self.chunk+= " "
	self.prefix=""

    def inc(self):
	self.level+= self.step
	self.prefix += self.chunk

    def dec(self):
	self.level-= self.step
	if(self.level < 0):
	    raise Exception, "Illegal indent level"
	self.prefix = self.prefix[:self.level]

    def __call__(self):
	return  self.prefix

class CFile(file):

    def __init__(self, name, flags):
        file.__init__(self, name, flags)
        self.indent = Indenter()
        self.newline = True

    def rwrite(self, value):
        file.write(self, value)

    def write(self, value=""):
        if self.newline:
            self.rwrite(self.indent())
            self.newline = False
        self.rwrite(value)

    def writeln(self, value=""):
        self.write(value+"\n")
        self.newline = True

    def beginBlock(self):
        self.writeln("{")
        self.indent.inc();

    def endBlock(self, suffix=""):
        self.indent.dec();
        self.writeln("}"+suffix)

CPP_PRIMITIVE_MAP = {
    "void" : "void",
    "bool" : "bool",
    "string": "std::string",
    "utf7": "std::string",
    "utf8": "std::wstring",
    "utf16": "std::utf16",
    "byte" : "uint8_t",
    "i08": "int8_t",
    "i16": "int16_t",
    "i32": "int32_t",
    "i64": "int64_t",
    "u08": "uint8_t",
    "u16": "uint16_t",
    "u32": "uint32_t",
    "u64": "uint64_t",
    "float": "double"
}

CPP_CONTAINER_MAP = {
    Map : "std::map",
    List: "std::list",
    Set : "std::set",
}

def typeToCTypeDeclaration(ttype):

    if isinstance(ttype, PrimitiveType):
        return CPP_PRIMITIVE_MAP[ttype.name]

    elif isinstance(ttype, CollectionType):

        result = CPP_CONTAINER_MAP[type(ttype)]+"<"
        
        if isinstance(ttype, Map):
            result+= typeToCTypeDeclaration(ttype.keyType)+", "+ typeToCTypeDeclaration(ttype.valueType)

        elif isinstance(ttype, Set) or isinstance(ttype, List):
            result+= typeToCTypeDeclaration(ttype.valueType)

        else:
            raise Exception, "Unknown Collection Type "+str(ttype)

        result+= "> "

        return result

    elif isinstance(ttype, Struct):
        return "struct "+ttype.name

    elif isinstance(ttype, TypeDef):
        return ttype.name;

    elif isinstance(ttype, Enum):
        return ttype.name;

    elif isinstance(ttype, Function):
        return typeToCTypeDeclaration(ttype.resultType)+ " "+ttype.name+"("+string.join([typeToCTypeDeclaration(arg) for arg in ttype.argFieldList], ", ")+")"

    elif isinstance(ttype, Field):
        return typeToCTypeDeclaration(ttype.type)+ " "+ttype.name

    else:
        raise Exception, "Unknown type "+str(ttype)

def writeTypeDefDefinition(cfile, typedef):

    cfile.writeln("typedef "+typeToCTypeDeclaration(typedef.definitionType)+" "+typedef.name+";")

def writeEnumDefinition(cfile, enum):

    cfile.write("enum "+enum.name+" ");

    cfile.beginBlock();

    first = True

    for ed in enum.enumDefs:
        if first:
            first = False
        else:
            cfile.writeln(",")
        cfile.write(ed.name+" = "+str(ed.id))

    cfile.writeln()
    cfile.endBlock(";");

def writeStructDefinition(cfile, struct):

    cfile.write("struct "+struct.name+" ");

    cfile.beginBlock()

    for field in struct.fieldList:
        cfile.writeln(typeToCTypeDeclaration(field)+";")

    cfile.endBlock(";")


CPP_DEFINITION_WRITER_MAP = {
    TypeDef : writeTypeDefDefinition,
    Enum : writeEnumDefinition,
    Struct : writeStructDefinition,
    Service : None
    }
    
def writeDefinitions(cfile, definitions):
    for definition in definitions:

        writer = CPP_DEFINITION_WRITER_MAP[type(definition)]

        if writer:
            writer(cfile, definition)

        cfile.writeln()

CPP_THRIFT_NS = "facebook::thrift"

CPP_INTERFACE_FUNCTION_DECLARATION = Template("""    virtual ${functionDeclaration} = 0;
""")

CPP_INTERFACE_DECLARATION = Template("""
class ${service}If {
    public:
    ~${service}If() {}
${functionDeclarations}};
""")

def writeServiceInterfaceDeclaration(cfile, service, debugp=None):

    functionDeclarations = string.join([CPP_INTERFACE_FUNCTION_DECLARATION.substitute(service=service.name, functionDeclaration=typeToCTypeDeclaration(function)) for function in service.functionList], "")

    cfile.write(CPP_INTERFACE_DECLARATION.substitute(service=service.name, functionDeclarations=functionDeclarations))

CPP_SP = Template("boost::shared_ptr<${klass}> ")

CPP_PROCESSOR = CPP_THRIFT_NS+"::TProcessor"
CPP_PROCESSORP = CPP_SP.substitute(klass=CPP_PROCESSOR)

CPP_PROTOCOL_NS = CPP_THRIFT_NS+"::protocol"
CPP_PROTOCOL = CPP_PROTOCOL_NS+"::TProcotol"
CPP_PROTOCOLP = CPP_SP.substitute(klass=CPP_PROTOCOL)


CPP_TRANSPORT_NS = CPP_THRIFT_NS+"::transport"
CPP_TRANSPORT = CPP_TRANSPORT_NS+"::TTransport"
CPP_TRANSPORTP = CPP_SP.substitute(klass=CPP_TRANSPORT)

CPP_SERVER_FUNCTION_DECLARATION = Template("""    void process_${function}("""+CPP_TRANSPORTP+""" _itrans, """+CPP_TRANSPORTP+""" _otrans);
""")

CPP_PROTOCOL_TSTOP = CPP_PROTOCOL_NS+"::T_STOP"
CPP_PROTOCOL_TTYPE = CPP_PROTOCOL_NS+"::TType"

CPP_SERVER_DECLARATION = Template("""
class ${service}ServerIf : public ${service}If, public """+CPP_PROCESSOR+""" {
    public:
    ${service}ServerIf("""+CPP_PROTOCOLP+""" protocol): _iprot(protocol), _oprot(protocol) {}
    ${service}ServerIf("""+CPP_PROTOCOLP+""" iprot, """+CPP_PROTOCOLP+""" oprot) : _iprot(iprot), _oprot(oprot) {}
    virtual ~${service}ServerIf() {}
    bool process("""+CPP_TRANSPORTP+""" _itrans,"""+CPP_TRANSPORTP+""" _otrans);
    protected:
    """+CPP_PROTOCOLP+""" _iprot;
    """+CPP_PROTOCOLP+""" _oprot;
    private:
${functionDeclarations}};
""")

def writeServerDeclaration(cfile, service, debugp=None):

    functionDeclarations = string.join([CPP_SERVER_FUNCTION_DECLARATION.substitute(function=function.name) for function in service.functionList], "")

    cfile.write(CPP_SERVER_DECLARATION.substitute(service=service.name, functionDeclarations=functionDeclarations))
    
CPP_CLIENT_FUNCTION_DECLARATION = Template("""    ${functionDeclaration};
""")

CPP_CLIENT_DECLARATION = Template("""
class ${service}Client : public ${service}If {

    public:

    ${service}Client("""+CPP_TRANSPORTP+""" transport, """+CPP_PROTOCOLP+""" protocol): _itrans(transport), _otrans(transport), _iprot(protocol), _oprot(protocol {}

    ${service}Client("""+CPP_TRANSPORTP+""" itrans, """+CPP_TRANSPORTP+""" otrans, """+CPP_PROTOCOLP+""" iprot, """+CPP_PROTOCOLP+""" oprot) : _itrans(itrans), _otrans(otrans), _iprot(iprot), _oprot(oprot)x {}

${functionDeclarations}};
""")

def writeClientDeclaration(cfile, service, debugp=None):

    functionDeclarations = string.join([CPP_CLIENT_FUNCTION_DECLARATION.substitute(functionDeclaration=typeToCTypeDeclaration(function)) for function in service.functionList], "")

    cfile.writeln(CPP_CLIENT_DECLARATION.substitute(service=service.name, functionDeclarations=functionDeclarations))

def writeServiceDeclaration(cfile, service, debugp=None):
    writeServiceInterfaceDeclaration(cfile, service, debugp)
    writeServerDeclaration(cfile, service, debugp)
    writeClientDeclaration(cfile, service, debugp)

def toGenDir(filename, suffix="cpp-gen", debugp=None):

    result = os.path.join(os.path.split(filename)[0], suffix)

    if not os.path.exists(result):
        os.mkdir(result)

    return result

def toBasename(filename, debugp=None):
    """ Take the filename minus the path and\".thrift\" extension  if present """

    basename = os.path.split(filename)[1]

    tokens = os.path.splitext(basename)

    if tokens[1].lower() == ".thrift":
        basename = tokens[0]

    if debugp:
        debugp("toBasename("+str(filename)+") => "+str(basename))

    return basename

def toDefinitionHeaderName(filename, genDir=None, debugp=None):

    if not genDir:
        genDir = toGenDir(filename)

    basename = toBasename(filename)

    result = os.path.join(genDir, basename+"_types.h")

    if debugp:
        debugp("toDefinitionHeaderName("+str(filename)+", "+str(genDir)+") => "+str(basename))

    return result

def writeDefinitionHeader(program, filename, genDir=None, debugp=None):

    definitionHeader = toDefinitionHeaderName(filename, genDir)

    if debugp:
        debugp("definitionHeader: "+str(definitionHeader))

    cfile = CFile(definitionHeader, "w")

    basename = toBasename(filename)

    cfile.writeln(CPP_TYPES_HEADER.substitute(source=basename, date=time.ctime()))

    writeDefinitions(cfile, program.definitions)

    cfile.writeln(CPP_TYPES_FOOTER.substitute(source=basename))

    cfile.close()

def toServicesHeaderName(filename, genDir=None, debugp=None):

    if not genDir:
        genDir = toGenDir(filename)

    basename = toBasename(filename)

    result = os.path.join(genDir, basename+".h")

    if debugp:
        debugp("toDefinitionHeaderName("+str(filename)+", "+str(genDir)+") => "+str(basename))

    return result


def writeServicesHeader(program, filename, genDir=None, debugp=None):

    servicesHeader = toServicesHeaderName(filename, genDir)

    if debugp:
        debugp("servicesHeader: "+str(servicesHeader))

    cfile = CFile(servicesHeader, "w")

    basename = toBasename(filename)

    cfile.writeln(CPP_SERVICES_HEADER.substitute(source=basename, date=time.ctime()))

    services = []

    # Build orderered list of service definitions by scanning definitions list for services

    for definition in  program.definitions:
        if isinstance(definition, Service) and definition.name in program.serviceMap:
            services.append(definition)

    for service in services:

        writeServiceDeclaration(cfile, service)

    cfile.writeln(CPP_SERVICES_FOOTER.substitute(source=basename))

    cfile.close()


CPP_STRUCT_READ = Template("""void read${name}Struct("""+CPP_PROTOCOLP+""" _iprot, """+CPP_TRANSPORTP+""" itrans, ${declaration}& value) {
    std::string name;
    uint32_t id;
    uint32_t type;
    while(true) {
        _iprot->readFieldBegin(_itrans, name, type, id);
        if(type == """+CPP_PROTOCOL_TSTOP+""") {
            break;
        }
        switch(id) {
${readFieldListSwitch}
        }
    }
}
""")

CPP_PRIMITIVE_TYPE_IO_METHOD_SUFFIX_MAP = {
    "bool" : "Bool",
    "string": "String",
    "utf7": "String",
    "utf8": "String", 
    "utf16": "String",
    "i08": "Byte",
    "i16": "I16",
    "i32": "I32", 
    "i64": "I64", 
    "u08": "Byte",
    "u16": "U16",
    "u32": "U32",
    "u64": "U64",
    "float": "Double"
}

CPP_COLLECTION_TYPE_IO_METHOD_SUFFIX_MAP = {
    Map : "stdmap",
    List : "stdlist",
    Set : "stdset"
}

def typeToIOMethodSuffix(ttype):

    if isinstance(ttype, PrimitiveType):
        return CPP_PRIMITIVE_TYPE_IO_METHOD_SUFFIX_MAP[ttype.name]

    elif isinstance(ttype, CollectionType):

        result = CPP_COLLECTION_TYPE_IO_METHOD_SUFFIX_MAP[type(ttype)]+"_"

        if isinstance(ttype, Map):
            result+= "k_"+typeToIOMethodSuffix(ttype.keyType)+"_"

        result += "v_"+typeToIOMethodSuffix(ttype.valueType)

        return result

    elif isinstance(ttype, Struct):
        return "struct_"+ttype.name

    elif isinstance(ttype, TypeDef):
        return typeToIOMethodSuffix(ttype.definitionType)

    elif isinstance(ttype, Enum):
        return typeToIOMethodSuffix(U32_TYPE)

    else:
        raise Exception, "Unknown type "+str(ttype)

def toReadCall(value, ttype):

    suffix = typeToIOMethodSuffix(ttype)

    if isinstance(ttype, PrimitiveType):
        return "iprot->read"+suffix+"(itrans, "+value+")"

    elif isinstance(ttype, CollectionType):
        return "read_"+suffix+"(iprot, itrans, "+value+")"

    elif isinstance(ttype, Struct):
        return "read_"+suffix+"(iprot, itrans, "+value+")"

    elif isinstance(ttype, TypeDef):
        return toReadCall(value, ttype.definitionType)

    elif isinstance(ttype, Enum):
        return toReadCall(value, U32_TYPE)

    else:
        raise Exception, "Unknown type "+str(ttype)

CPP_READ_MAP_DEFINITION = Template("""
void read_${suffix}("""+CPP_PROTOCOLP+""" iprot, """+CPP_TRANSPORTP+""" itrans, ${declaration}& value) {

   uint32_t count;
   ${keyType} key;
   ${valueType} elem;

   _iprot->readU32(itrans, count);

   for(int ix = 0; ix <  count; ix++) {
       ${keyReadCall};
       ${valueReadCall};
       value.insert(std::make_pair(key, elem));
   }
}
""")
    
CPP_WRITE_MAP_DEFINITION = Template("""
void write_${suffix}("""+CPP_PROTOCOLP+""" oprot, """+CPP_TRANSPORTP+""" otrans, ${declaration}& value) {

   uint32_t count;
   ${keyType} key;
   ${valueType} elem;

   _oprot->writeU32(otrans, count);

   for(int ix = 0; ix <  count; ix++) {
       ${keyReadCall};
       ${valueReadCall};
       value.insert(std::make_pair(key, elem));
   }
}
""")
    

CPP_READ_LIST_DEFINITION = Template("""
void read_${suffix}("""+CPP_PROTOCOLP+""" iprot, """+CPP_TRANSPORTP+""" itrans, ${declaration}& value) {

   uint32_t count;
   ${valueType} elem;

   _iprot->readU32(itrans,  count);

   for(int ix = 0; ix < count; ix++) {
       ${valueReadCall};
       value.insert(elem);
   }
}
""")
    
def toCollectionReadDefinition(ttype):

    suffix = typeToIOMethodSuffix(ttype)

    if isinstance(ttype, Map):
        keyReadCall = toReadCall("key", ttype.keyType)

    valueReadCall= toReadCall("elem", ttype.valueType)

    if isinstance(ttype, Map):
        return CPP_READ_MAP_DEFINITION.substitute(suffix=suffix, declaration=typeToCTypeDeclaration(ttype),
                                                  keyType=typeToCTypeDeclaration(ttype.keyType),
                                                  keyReadCall=keyReadCall,
                                                  valueType=typeToCTypeDeclaration(ttype.valueType),
                                                  valueReadCall=valueReadCall)

    else:
        return CPP_READ_LIST_DEFINITION.substitute(suffix=suffix, declaration=typeToCTypeDeclaration(ttype),
                                                   valueReadCall=valueReadCall,
                                                   valueType=typeToCTypeDeclaration(ttype.valueType))


CPP_READ_STRUCT_DEFINITION = Template("""
void read_${suffix}("""+CPP_PROTOCOLP+""" iprot, """+CPP_TRANSPORTP+""" itrans, ${declaration}& value) {

    std::string name;
    """+CPP_PROTOCOL_TTYPE+""" type;
    uint16_t id;

    while(true) {

        _iprot->readFieldBegin(itrans, name, type, id);

        if(type == """+CPP_PROTOCOL_TSTOP+""") {break;}

        switch(id) {
${fieldSwitch}
            default:
            iprot->skip(itrans, type);
            break;
        }

        _iprot->readFieldEnd(itrans);
    }
}
""")
    
def toStructReadDefinition(ttype):

    suffix = typeToIOMethodSuffix(ttype)

    # Sort field list in order of increasing ids

    fieldList = []
    fieldList+= ttype.fieldList

    fieldList.sort(lambda a,b: a.id - b.id)

    fieldSwitch=""

    for field in fieldList:
        fieldSwitch+= "            case "+str(field.id)+": "
        fieldSwitch+= toReadCall("value."+field.name, field.type)+"; break;\n"

    return CPP_READ_STRUCT_DEFINITION.substitute(suffix=suffix, declaration=typeToCTypeDeclaration(ttype), fieldSwitch=fieldSwitch)
    
def toReadDefinition(ttype):
    if isinstance(ttype, CollectionType):
        return toCollectionReadDefinition(ttype)

    elif isinstance(ttype, Struct):
        return toStructReadDefinition(ttype)

class CPPGenerator(Generator):

    def __call__(self, program, filename, genDir=None, debugp=None):

        writeDefinitionHeader(program, filename, gendir, debugp)
        
        writeServicesHeader(program, filename, gendir, debugp)
        
        writeClientHeader(program, filename, gendir, debugp)