<commit_before>#include <cstdlib>
#include <iostream>
#include <vector>
#include <string>
#include <stdlib.h>
#include <sstream>
#include <climits>
using namespace std;

// Command line parsing
#include <getopt.h>

// DyninstAPI includes
#include "BPatch.h"
#include "BPatch_binaryEdit.h"
#include "BPatch_flowGraph.h"
#include "BPatch_function.h"
#include "BPatch_point.h"

using namespace Dyninst;

//cmd line options
char *originalBinary;
char *instrumentedBinary;
bool verbose = false;
Dyninst::Address entryPoint;
set<string> instrumentLibraries;
set<string> runtimeLibraries;
int bbSkip = 0;
bool skipMainModule = false;

const char *instLibrary = "libAflDyninst.so";

static const char *OPT_STR = "i:o:l:e:vs:dr:";
static const char *USAGE = " -i <binary> -o <binary> -l <library> -e <address> -s <number>\n \
            -i: Input binary \n \
            -o: Output binary\n \
            -d: Don't instrument the binary, only supplied libraries\n \
            -l: Linked library to instrument (repeat for more than one)\n \
            -r: Runtime library to instrument (path to, repeat for more than one)\n \
            -e: Entry point address to patch (required for stripped binaries)\n \
            -s: Number of basic blocks to skip\n \
            -v: Verbose output\n";

bool parseOptions(int argc, char **argv)
{

    int c;
    while ((c = getopt (argc, argv, OPT_STR)) != -1) {
        switch ((char) c) {
        case 'e':
            entryPoint = strtoul(optarg, NULL, 16);;
            break;
        case 'i':
            originalBinary = optarg;
            instrumentLibraries.insert(optarg);
            break;
        case 'o':
            instrumentedBinary = optarg;
            break;
        case 'l':
            instrumentLibraries.insert(optarg);
            break;
        case 'r':
            runtimeLibraries.insert(optarg);
            break;
        case 's':
            bbSkip = atoi(optarg);
            break;
        case 'd':
            skipMainModule = true;
            break;
        case 'v':
            verbose = true;
            break;
        default:
            cerr << "Usage: " << argv[0] << USAGE;
            return false;
        }
    }

    if(originalBinary == NULL) {
        cerr << "Input binary is required!"<< endl;
        cerr << "Usage: " << argv[0] << USAGE;
        return false;
    }

    if(instrumentedBinary == NULL) {
        cerr << "Output binary is required!" << endl;
        cerr << "Usage: " << argv[0] << USAGE;
        return false;
    }

    if(skipMainModule && instrumentLibraries.empty()) {
        cerr << "If using option -d , option -l is required." << endl;
        cerr << "Usage: " << argv[0] << USAGE;
        return false;
    }

    return true;
}

BPatch_function *findFuncByName (BPatch_image * appImage, char *funcName)
{
    BPatch_Vector < BPatch_function * >funcs;

    if (NULL == appImage->findFunction (funcName, funcs) || !funcs.size ()
        || NULL == funcs[0]) {
        cerr << "Failed to find " << funcName << " function." << endl;
        return NULL;
    }

    return funcs[0];
}

// insert callback to initialization function in the instrumentation library
// either at _init or at manualy specified entry point.
bool insertCallToInit(BPatch_binaryEdit * appBin, BPatch_function * instIncFunc,BPatch_module * module, BPatch_function *funcInit)
{

    /* Find the instrumentation points */
    vector<BPatch_point *> points;
    vector < BPatch_point * >*funcEntry = funcInit->findPoint (BPatch_entry);

    if (NULL == funcEntry) {
        cerr << "Failed to find entry for function. " <<  endl;
        return false;
    }

    cout << "Inserting init callback." << endl;
    BPatch_Vector < BPatch_snippet * >instArgs; // init has no args
    BPatch_funcCallExpr instIncExpr (*instIncFunc, instArgs);

    /* Insert the snippet at function entry */
    BPatchSnippetHandle *handle =
        appBin->insertSnippet (instIncExpr, *funcEntry, BPatch_callBefore,
                               BPatch_lastSnippet);
    if (!handle) {
        cerr << "Failed to insert init callback." << endl;
        return false;
    }
    return true;
}

// inserts a callback for each basic block assigning it an instrumentation
// time 16bit random ID just as afl
bool insertBBCallback(BPatch_binaryEdit * appBin, BPatch_function * curFunc,
                      char *funcName, BPatch_function * instBBIncFunc,int *bbIndex)
{

    BPatch_flowGraph *appCFG = curFunc->getCFG ();
    unsigned short randID;

    if (!appCFG) {
        cerr << "Failed to find CFG for function " << funcName << endl;
        return false;
    }

    BPatch_Set < BPatch_basicBlock * >allBlocks;
    if (!appCFG->getAllBasicBlocks (allBlocks)) {
        cerr << "Failed to find basic blocks for function " << funcName << endl;
        return false;
    } else if (allBlocks.size () == 0) {
        cerr << "No basic blocks for function " << funcName << endl;
        return false;
    }

    BPatch_Set < BPatch_basicBlock * >::iterator iter;
    for (iter = allBlocks.begin (); iter != allBlocks.end (); iter++) {
        if(*bbIndex < bbSkip) { // skip over first bbSkip bbs
            (*bbIndex)++;
            continue;

        }
        unsigned long address = (*iter)->getStartAddress ();
        randID = rand() % USHRT_MAX;
        if(verbose) {
            cout << "Instrumenting Basic Block 0x" << hex << address << " of " <<
                 funcName << " with random id " << randID << endl;
        }

        BPatch_Vector < BPatch_snippet * >instArgs;
        BPatch_constExpr bbId (randID);
        instArgs.push_back (&bbId);
        BPatch_point *bbEntry = (*iter)->findEntryPoint ();
        if (NULL == bbEntry) {
            // warn the user, but continue
            cerr << "Failed to find entry for basic block at 0x" << hex << address
                 << endl;
            (*bbIndex)++;
            continue;
        }
        BPatch_funcCallExpr instIncExpr (*instBBIncFunc, instArgs);
        BPatchSnippetHandle *handle =
            appBin->insertSnippet (instIncExpr, *bbEntry, BPatch_callBefore,
                                   BPatch_lastSnippet);
        if (!handle) {
            // warn the user, but continue to next bb
            cerr << "Failed to insert instrumention in basic block at 0x" << hex <<
                 address << endl;
            (*bbIndex)++;
            continue;
        }

        (*bbIndex)++;
    }

    return true;

}

int main (int argc, char **argv)
{

    if(!parseOptions(argc,argv)) {
        return EXIT_FAILURE;
    }

    BPatch bpatch;

    BPatch_binaryEdit *appBin = bpatch.openBinary (originalBinary, instrumentLibraries.size() != 1);
    if (appBin == NULL) {
        cerr << "Failed to open binary" << endl;
        return EXIT_FAILURE;
    }

    BPatch_image *appImage = appBin->getImage ();

    //get and iterate over all modules, instrumenting only the default and manualy specified ones
    vector < BPatch_module * >*modules = appImage->getModules ();
    vector < BPatch_module * >::iterator moduleIter;
    vector < BPatch_function * >* funcsInModule;
    BPatch_module *defaultModule = NULL;
    string defaultModuleName;

    // look for _init
    if(defaultModuleName.empty()){
        for (moduleIter = modules->begin(); moduleIter != modules->end(); ++moduleIter){
            funcsInModule = (*moduleIter)->getProcedures();
    	    vector < BPatch_function * >::iterator funcsIterator;
    	    for (funcsIterator = funcsInModule->begin(); funcsIterator != funcsInModule->end(); ++funcsIterator){
        		char funcName[1024];
        		(*funcsIterator)->getName(funcName,1024);
        		if(string(funcName) == string("_init")){
    	            char moduleName[1024];
        		    (*moduleIter)->getName(moduleName,1024);
    	            defaultModuleName = string(moduleName);
                    if(verbose) {
                        cout << "Found _init in " << moduleName<< endl;
                    }
        		    break;
        		}
            }
                    if(!defaultModuleName.empty()) break;
    	}
    }
    // last resort, by name of the binary
    if(defaultModuleName.empty())
        defaultModuleName = string(originalBinary).substr(string(originalBinary).find_last_of("\\/")+1);

    if (!appBin->loadLibrary (instLibrary)) {
        cerr << "Failed to open instrumentation library." << endl;
        cerr << "It needs to be located in the current working directory." << endl;
        return EXIT_FAILURE;
    }

	appImage = appBin->getImage ();

    /* Find code coverage functions in the instrumentation library */
    BPatch_function *initAflForkServer =
        findFuncByName (appImage, (char *) "initAflForkServer");
    BPatch_function *bbCallback =
        findFuncByName (appImage, (char *) "bbCallback");
    if (!initAflForkServer || !bbCallback ) {
        cerr << "Instrumentation library lacks callbacks!" << endl;
        return EXIT_FAILURE;
    }

    int bbIndex = 0;
    for (moduleIter = modules->begin (); moduleIter != modules->end (); ++moduleIter) {
        char moduleName[1024];
        (*moduleIter)->getName (moduleName, 1024);

        if ((*moduleIter)->isSharedLib ()) {
            if (instrumentLibraries.find (moduleName) == instrumentLibraries.end ()) {
                cout << "Skipping library: " << moduleName << endl;
                continue;
            }
        }

        if (string (moduleName).find (defaultModuleName) != string::npos) {
            defaultModule = (*moduleIter);
            if(skipMainModule) continue;
        }
        cout << "Instrumenting module: " << moduleName << endl;
        vector < BPatch_function * >*allFunctions =
            (*moduleIter)->getProcedures ();
        vector < BPatch_function * >::iterator funcIter;

        // iterate over all functions in the module
        for (funcIter = allFunctions->begin (); funcIter != allFunctions->end ();
             ++funcIter) {
            BPatch_function *curFunc = *funcIter;
            char funcName[1024];
            curFunc->getName (funcName, 1024);
            if(string (funcName) == string("_start")) continue; // here's a bug on hlt
            insertBBCallback (appBin, curFunc, funcName, bbCallback, &bbIndex);
        }

    }

    //if entrypoint set ,find function  , else find _init
    BPatch_function *funcToPatch = NULL;
    if(!entryPoint) {
        BPatch_Vector<BPatch_function*> funcs;
        defaultModule->findFunction("_init", funcs);
        if(!funcs.size()) {
            cerr << "Couldn't locate _init, specify entry point manualy. "<< endl;
            return EXIT_FAILURE;
        }
        // there should really be only one
        funcToPatch = funcs[0];
    } else {
        funcToPatch = defaultModule->findFunctionByEntry(entryPoint);
    }
    if(!funcToPatch) {
        cerr << "Couldn't locate function at given entry point. "<< endl;
        return EXIT_FAILURE;
    }
    if(!insertCallToInit (appBin,  initAflForkServer,defaultModule,funcToPatch)){
        cerr << "Could not insert init callback at given entry point." << endl;
        return EXIT_FAILURE;
    }

    cout << "Saving the instrumented binary to " << instrumentedBinary << "..." << endl;
    // Output the instrumented binary
    if (!appBin->writeFile (instrumentedBinary)) {
        cerr << "Failed to write output file: " << instrumentedBinary << endl;
        return EXIT_FAILURE;
    }

    if(!runtimeLibraries.empty()) {
        cout << "Instrumenting runtime libraries." << endl;
        set<string>::iterator rtLibIter ;
        for(rtLibIter = runtimeLibraries.begin(); rtLibIter != runtimeLibraries.end(); rtLibIter++) {
            BPatch_binaryEdit *libBin = bpatch.openBinary ((*rtLibIter).c_str(), false);
            printf("I sad otvara %s\n",(*rtLibIter).c_str());
            if (libBin == NULL) {
                cerr << "Failed to open binary "<< *rtLibIter << endl;
                return EXIT_FAILURE;
            }
            BPatch_image *libImg = libBin->getImage ();
            vector < BPatch_module * >*modules = libImg->getModules ();
            libBin->loadLibrary (instLibrary);
            moduleIter = modules->begin ();

            for ( ; moduleIter != modules->end (); ++moduleIter) {
                char moduleName[1024];
                (*moduleIter)->getName (moduleName, 1024);
                cout << "Instrumenting module: " << moduleName << endl;
                vector < BPatch_function * >*allFunctions =
                    (*moduleIter)->getProcedures ();
                vector < BPatch_function * >::iterator funcIter;
                // iterate over all functions in the module
                for (funcIter = allFunctions->begin (); funcIter != allFunctions->end ();
                     ++funcIter) {
                    BPatch_function *curFunc = *funcIter;
                    char funcName[1024];
                    curFunc->getName (funcName, 1024);
                    if(string (funcName) == string("_start")) continue;
                    insertBBCallback (libBin, curFunc, funcName, bbCallback, &bbIndex);
                }
            }
            if (!libBin->writeFile ((*rtLibIter + ".ins").c_str())) {
                cerr << "Failed to write output file: " <<(*rtLibIter + ".ins").c_str() << endl;
                return EXIT_FAILURE;
            } else {
                cout << "Saved the instrumented library to " << (*rtLibIter + ".ins").c_str() << "." << endl;
            }
        }
    }

    cout << "All done! Happy fuzzing!" << endl;
    return EXIT_SUCCESS;

}
<commit_msg>removed superfluo superfluous lines<commit_after>#include <cstdlib>
#include <iostream>
#include <vector>
#include <string>
#include <stdlib.h>
#include <sstream>
#include <climits>
using namespace std;

// Command line parsing
#include <getopt.h>

// DyninstAPI includes
#include "BPatch.h"
#include "BPatch_binaryEdit.h"
#include "BPatch_flowGraph.h"
#include "BPatch_function.h"
#include "BPatch_point.h"

using namespace Dyninst;

//cmd line options
char *originalBinary;
char *instrumentedBinary;
bool verbose = false;
Dyninst::Address entryPoint;
set<string> instrumentLibraries;
set<string> runtimeLibraries;
int bbSkip = 0;
bool skipMainModule = false;

const char *instLibrary = "libAflDyninst.so";

static const char *OPT_STR = "i:o:l:e:vs:dr:";
static const char *USAGE = " -i <binary> -o <binary> -l <library> -e <address> -s <number>\n \
            -i: Input binary \n \
            -o: Output binary\n \
            -d: Don't instrument the binary, only supplied libraries\n \
            -l: Linked library to instrument (repeat for more than one)\n \
            -r: Runtime library to instrument (path to, repeat for more than one)\n \
            -e: Entry point address to patch (required for stripped binaries)\n \
            -s: Number of basic blocks to skip\n \
            -v: Verbose output\n";

bool parseOptions(int argc, char **argv)
{

    int c;
    while ((c = getopt (argc, argv, OPT_STR)) != -1) {
        switch ((char) c) {
        case 'e':
            entryPoint = strtoul(optarg, NULL, 16);;
            break;
        case 'i':
            originalBinary = optarg;
            instrumentLibraries.insert(optarg);
            break;
        case 'o':
            instrumentedBinary = optarg;
            break;
        case 'l':
            instrumentLibraries.insert(optarg);
            break;
        case 'r':
            runtimeLibraries.insert(optarg);
            break;
        case 's':
            bbSkip = atoi(optarg);
            break;
        case 'd':
            skipMainModule = true;
            break;
        case 'v':
            verbose = true;
            break;
        default:
            cerr << "Usage: " << argv[0] << USAGE;
            return false;
        }
    }

    if(originalBinary == NULL) {
        cerr << "Input binary is required!"<< endl;
        cerr << "Usage: " << argv[0] << USAGE;
        return false;
    }

    if(instrumentedBinary == NULL) {
        cerr << "Output binary is required!" << endl;
        cerr << "Usage: " << argv[0] << USAGE;
        return false;
    }

    if(skipMainModule && instrumentLibraries.empty()) {
        cerr << "If using option -d , option -l is required." << endl;
        cerr << "Usage: " << argv[0] << USAGE;
        return false;
    }

    return true;
}

BPatch_function *findFuncByName (BPatch_image * appImage, char *funcName)
{
    BPatch_Vector < BPatch_function * >funcs;

    if (NULL == appImage->findFunction (funcName, funcs) || !funcs.size ()
        || NULL == funcs[0]) {
        cerr << "Failed to find " << funcName << " function." << endl;
        return NULL;
    }

    return funcs[0];
}

// insert callback to initialization function in the instrumentation library
// either at _init or at manualy specified entry point.
bool insertCallToInit(BPatch_binaryEdit * appBin, BPatch_function * instIncFunc,BPatch_module * module, BPatch_function *funcInit)
{

    /* Find the instrumentation points */
    vector<BPatch_point *> points;
    vector < BPatch_point * >*funcEntry = funcInit->findPoint (BPatch_entry);

    if (NULL == funcEntry) {
        cerr << "Failed to find entry for function. " <<  endl;
        return false;
    }

    cout << "Inserting init callback." << endl;
    BPatch_Vector < BPatch_snippet * >instArgs; // init has no args
    BPatch_funcCallExpr instIncExpr (*instIncFunc, instArgs);

    /* Insert the snippet at function entry */
    BPatchSnippetHandle *handle =
        appBin->insertSnippet (instIncExpr, *funcEntry, BPatch_callBefore,
                               BPatch_lastSnippet);
    if (!handle) {
        cerr << "Failed to insert init callback." << endl;
        return false;
    }
    return true;
}

// inserts a callback for each basic block assigning it an instrumentation
// time 16bit random ID just as afl
bool insertBBCallback(BPatch_binaryEdit * appBin, BPatch_function * curFunc,
                      char *funcName, BPatch_function * instBBIncFunc,int *bbIndex)
{

    BPatch_flowGraph *appCFG = curFunc->getCFG ();
    unsigned short randID;

    if (!appCFG) {
        cerr << "Failed to find CFG for function " << funcName << endl;
        return false;
    }

    BPatch_Set < BPatch_basicBlock * >allBlocks;
    if (!appCFG->getAllBasicBlocks (allBlocks)) {
        cerr << "Failed to find basic blocks for function " << funcName << endl;
        return false;
    } else if (allBlocks.size () == 0) {
        cerr << "No basic blocks for function " << funcName << endl;
        return false;
    }

    BPatch_Set < BPatch_basicBlock * >::iterator iter;
    for (iter = allBlocks.begin (); iter != allBlocks.end (); iter++) {
        if(*bbIndex < bbSkip) { // skip over first bbSkip bbs
            (*bbIndex)++;
            continue;

        }
        unsigned long address = (*iter)->getStartAddress ();
        randID = rand() % USHRT_MAX;
        if(verbose) {
            cout << "Instrumenting Basic Block 0x" << hex << address << " of " <<
                 funcName << " with random id " << randID << endl;
        }

        BPatch_Vector < BPatch_snippet * >instArgs;
        BPatch_constExpr bbId (randID);
        instArgs.push_back (&bbId);
        BPatch_point *bbEntry = (*iter)->findEntryPoint ();
        if (NULL == bbEntry) {
            // warn the user, but continue
            cerr << "Failed to find entry for basic block at 0x" << hex << address
                 << endl;
            (*bbIndex)++;
            continue;
        }
        BPatch_funcCallExpr instIncExpr (*instBBIncFunc, instArgs);
        BPatchSnippetHandle *handle =
            appBin->insertSnippet (instIncExpr, *bbEntry, BPatch_callBefore,
                                   BPatch_lastSnippet);
        if (!handle) {
            // warn the user, but continue to next bb
            cerr << "Failed to insert instrumention in basic block at 0x" << hex <<
                 address << endl;
            (*bbIndex)++;
            continue;
        }

        (*bbIndex)++;
    }

    return true;

}

int main (int argc, char **argv)
{

    if(!parseOptions(argc,argv)) {
        return EXIT_FAILURE;
    }

    BPatch bpatch;

    BPatch_binaryEdit *appBin = bpatch.openBinary (originalBinary, instrumentLibraries.size() != 1);
    if (appBin == NULL) {
        cerr << "Failed to open binary" << endl;
        return EXIT_FAILURE;
    }

    BPatch_image *appImage = appBin->getImage ();

    //get and iterate over all modules, instrumenting only the default and manualy specified ones
    vector < BPatch_module * >*modules = appImage->getModules ();
    vector < BPatch_module * >::iterator moduleIter;
    vector < BPatch_function * >* funcsInModule;
    BPatch_module *defaultModule = NULL;
    string defaultModuleName;

    // look for _init
    if(defaultModuleName.empty()){
        for (moduleIter = modules->begin(); moduleIter != modules->end(); ++moduleIter){
            funcsInModule = (*moduleIter)->getProcedures();
    	    vector < BPatch_function * >::iterator funcsIterator;
    	    for (funcsIterator = funcsInModule->begin(); funcsIterator != funcsInModule->end(); ++funcsIterator){
        		char funcName[1024];
        		(*funcsIterator)->getName(funcName,1024);
        		if(string(funcName) == string("_init")){
    	            char moduleName[1024];
        		    (*moduleIter)->getName(moduleName,1024);
    	            defaultModuleName = string(moduleName);
                    if(verbose) {
                        cout << "Found _init in " << moduleName<< endl;
                    }
        		    break;
        		}
            }
                    if(!defaultModuleName.empty()) break;
    	}
    }
    // last resort, by name of the binary
    if(defaultModuleName.empty())
        defaultModuleName = string(originalBinary).substr(string(originalBinary).find_last_of("\\/")+1);

    if (!appBin->loadLibrary (instLibrary)) {
        cerr << "Failed to open instrumentation library." << endl;
        cerr << "It needs to be located in the current working directory." << endl;
        return EXIT_FAILURE;
    }

	appImage = appBin->getImage ();

    /* Find code coverage functions in the instrumentation library */
    BPatch_function *initAflForkServer =
        findFuncByName (appImage, (char *) "initAflForkServer");
    BPatch_function *bbCallback =
        findFuncByName (appImage, (char *) "bbCallback");
    if (!initAflForkServer || !bbCallback ) {
        cerr << "Instrumentation library lacks callbacks!" << endl;
        return EXIT_FAILURE;
    }

    int bbIndex = 0;
    for (moduleIter = modules->begin (); moduleIter != modules->end (); ++moduleIter) {
        char moduleName[1024];
        (*moduleIter)->getName (moduleName, 1024);

        if ((*moduleIter)->isSharedLib ()) {
            if (instrumentLibraries.find (moduleName) == instrumentLibraries.end ()) {
                cout << "Skipping library: " << moduleName << endl;
                continue;
            }
        }

        if (string (moduleName).find (defaultModuleName) != string::npos) {
            defaultModule = (*moduleIter);
            if(skipMainModule) continue;
        }
        cout << "Instrumenting module: " << moduleName << endl;
        vector < BPatch_function * >*allFunctions =
            (*moduleIter)->getProcedures ();
        vector < BPatch_function * >::iterator funcIter;

        // iterate over all functions in the module
        for (funcIter = allFunctions->begin (); funcIter != allFunctions->end ();
             ++funcIter) {
            BPatch_function *curFunc = *funcIter;
            char funcName[1024];
            curFunc->getName (funcName, 1024);
            if(string (funcName) == string("_start")) continue; // here's a bug on hlt
            insertBBCallback (appBin, curFunc, funcName, bbCallback, &bbIndex);
        }

    }

    //if entrypoint set ,find function  , else find _init
    BPatch_function *funcToPatch = NULL;
    if(!entryPoint) {
        BPatch_Vector<BPatch_function*> funcs;
        defaultModule->findFunction("_init", funcs);
        if(!funcs.size()) {
            cerr << "Couldn't locate _init, specify entry point manualy. "<< endl;
            return EXIT_FAILURE;
        }
        // there should really be only one
        funcToPatch = funcs[0];
    } else {
        funcToPatch = defaultModule->findFunctionByEntry(entryPoint);
    }
    if(!funcToPatch) {
        cerr << "Couldn't locate function at given entry point. "<< endl;
        return EXIT_FAILURE;
    }
    if(!insertCallToInit (appBin,  initAflForkServer,defaultModule,funcToPatch)){
        cerr << "Could not insert init callback at given entry point." << endl;
        return EXIT_FAILURE;
    }

    cout << "Saving the instrumented binary to " << instrumentedBinary << "..." << endl;
    // Output the instrumented binary
    if (!appBin->writeFile (instrumentedBinary)) {
        cerr << "Failed to write output file: " << instrumentedBinary << endl;
        return EXIT_FAILURE;
    }

    if(!runtimeLibraries.empty()) {
        cout << "Instrumenting runtime libraries." << endl;
        set<string>::iterator rtLibIter ;
        for(rtLibIter = runtimeLibraries.begin(); rtLibIter != runtimeLibraries.end(); rtLibIter++) {
            BPatch_binaryEdit *libBin = bpatch.openBinary ((*rtLibIter).c_str(), false);
            if (libBin == NULL) {
                cerr << "Failed to open binary "<< *rtLibIter << endl;
                return EXIT_FAILURE;
            }
            BPatch_image *libImg = libBin->getImage ();
            vector < BPatch_module * >*modules = libImg->getModules ();
            moduleIter = modules->begin ();

            for ( ; moduleIter != modules->end (); ++moduleIter) {
                char moduleName[1024];
                (*moduleIter)->getName (moduleName, 1024);
                cout << "Instrumenting module: " << moduleName << endl;
                vector < BPatch_function * >*allFunctions =
                    (*moduleIter)->getProcedures ();
                vector < BPatch_function * >::iterator funcIter;
                // iterate over all functions in the module
                for (funcIter = allFunctions->begin (); funcIter != allFunctions->end ();
                     ++funcIter) {
                    BPatch_function *curFunc = *funcIter;
                    char funcName[1024];
                    curFunc->getName (funcName, 1024);
                    if(string (funcName) == string("_start")) continue;
                    insertBBCallback (libBin, curFunc, funcName, bbCallback, &bbIndex);
                }
            }
            if (!libBin->writeFile ((*rtLibIter + ".ins").c_str())) {
                cerr << "Failed to write output file: " <<(*rtLibIter + ".ins").c_str() << endl;
                return EXIT_FAILURE;
            } else {
                cout << "Saved the instrumented library to " << (*rtLibIter + ".ins").c_str() << "." << endl;
            }
        }
    }

    cout << "All done! Happy fuzzing!" << endl;
    return EXIT_SUCCESS;

}
<|endoftext|>


<commit_before>/*=========================================================================

  Program:   Visualization Toolkit
  Module:    vtkOutputPort.cxx

  Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
  All rights reserved.
  See Copyright.txt or http://www.kitware.com/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notice for more information.

=========================================================================*/
#include "vtkOutputPort.h"

#include "vtkCommand.h"
#include "vtkDataObject.h"
#include "vtkDemandDrivenPipeline.h"
#include "vtkInformation.h"
#include "vtkInputPort.h"
#include "vtkMultiProcessController.h"
#include "vtkMultiProcessController.h"
#include "vtkObjectFactory.h"

vtkCxxRevisionMacro(vtkOutputPort, "1.11");
vtkStandardNewMacro(vtkOutputPort);

vtkCxxSetObjectMacro(vtkOutputPort,Controller,vtkMultiProcessController);

//----------------------------------------------------------------------------
vtkOutputPort::vtkOutputPort()
{
  this->Tag = -1;

  // Controller keeps a reference to this object as well.
  this->Controller = NULL;
  this->SetController(vtkMultiProcessController::GetGlobalController());

  this->PipelineFlag = 0;
  this->ParameterMethod = NULL;
  this->ParameterMethodArgDelete = NULL;
  this->ParameterMethodArg = NULL;

  this->SetNumberOfInputPorts(1);
}

//----------------------------------------------------------------------------
// We need to have a "GetNetReferenceCount" to avoid memory leaks.
vtkOutputPort::~vtkOutputPort()
{
  this->SetController(0);

  if ((this->ParameterMethodArg)&&(this->ParameterMethodArgDelete))
    {
    (*this->ParameterMethodArgDelete)(this->ParameterMethodArg);
    }
}

//----------------------------------------------------------------------------
void vtkOutputPort::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os,indent);
  os << indent << "Tag: " << this->Tag << endl;
  os << indent << "Controller: (" << this->Controller << ")\n";
  os << indent << "Pipeline Flag: "
     << (this->PipelineFlag ? "On\n" : "Off\n");
}

//----------------------------------------------------------------------------
// Remote method call to UpdateInformation and send the information downstream.
// This should be a friend.
void vtkOutputPortUpdateInformationCallBack(void *arg, void *remoteArgs,
                                    int remoteArgsLength, int remoteProcessId)
{
  vtkOutputPort *self = (vtkOutputPort*)arg;

  remoteArgs = remoteArgs;
  remoteArgsLength = remoteArgsLength;
  // Just call a method
  self->TriggerUpdateInformation(remoteProcessId);
}
//----------------------------------------------------------------------------
void vtkOutputPort::TriggerUpdateInformation(int remoteProcessId)
{
  vtkDataObject *input = this->GetInput();

  // Handle no input gracefully.
  if ( input != NULL )
    {
    input->UpdateInformation();
    }

  // The MTime of the input should also be considered.
  // Important for pipeline parallelism.
  // Include it in the information for efficiency.
  unsigned long latestMTime;
  latestMTime = input->GetMTime();
  vtkDemandDrivenPipeline *ddp =
    vtkDemandDrivenPipeline::SafeDownCast(this->GetExecutive());
  if (ddp)
    {
    ddp->UpdateInformation();
    unsigned long t2 = ddp->GetPipelineMTime();
    if (t2 > latestMTime)
      {
      latestMTime = t2;
      }
    }

  // Now just send the information downstream.
  // PipelineMTime is part of information, so downstream
  // port will make the time comparison, and call Update if necessary.
  int wholeInformation[8];
  input->GetWholeExtent( wholeInformation );

  this->Controller->Send( wholeInformation, 7,
                          remoteProcessId, vtkInputPort::INFORMATION_TRANSFER_TAG);

  this->Controller->Send( &latestMTime, 1,
                          remoteProcessId, vtkInputPort::INFORMATION_TRANSFER_TAG );

  int maxNumPieces = input->GetMaximumNumberOfPieces();
  this->Controller->Send( &maxNumPieces, 1,
                          remoteProcessId, vtkInputPort::INFORMATION_TRANSFER_TAG );
}

//----------------------------------------------------------------------------
// Remote method call to Update and send data downstream.
// This should be a friend.
void vtkOutputPortUpdateCallBack(void *arg, void *remoteArgs,
                                 int remoteArgsLength, int remoteProcessId)
{
  vtkOutputPort *self = (vtkOutputPort*)arg;

  remoteArgs = remoteArgs;
  remoteArgsLength = remoteArgsLength;
  // Just call a method
  self->TriggerUpdate(remoteProcessId);
}
//----------------------------------------------------------------------------
void vtkOutputPort::TriggerUpdate(int remoteProcessId)
{
  unsigned long downDataTime;
  vtkDataObject *input = this->GetInput();

  // First get the update extent requested.
  int extent[9];
  this->Controller->Receive( extent, 9, remoteProcessId,
                            vtkInputPort::UPDATE_EXTENT_TAG);
  input->SetUpdateExtent( extent );
  input->SetUpdatePiece( extent[6] );
  input->SetUpdateNumberOfPieces( extent[7] );
  input->SetUpdateGhostLevel( extent[8] );

  // Note:  Receiving DataTime was the start of a more intelligent promotion
  // for pipeline parallism.  Unfortunately there was no way (I knew of)
  // for us to not promote on the first Update.  I backed off, and am
  // requiring either 1: Filters handle Update Not returning correct data,
  // or 2: Pipeline parallelism must be primed with Updates on the Output
  // ports (giving correct data).  This Receive can be removed.

  // This is for pipeline parallism.
  // This Output port may or may not promote our data (execute).
  // We need the data time of the last transfer to compare to the mtime
  // of our input to determine if it should send the data (execute).
  this->Controller->Receive( &(downDataTime), 1, remoteProcessId,
                             vtkInputPort::NEW_DATA_TIME_TAG);

  // What was the idea of using relesed data here?  It caused a bug for multiple updates.
  // if ( input != NULL && input->GetDataReleased())

  // Postpone the update if we want pipeline parallism.
  // Handle no input gracefully. (Not true: Later we will send a NULL input.)
  if ( input != NULL && this->PipelineFlag == 0)
    {
    input->UpdateInformation();
    input->PropagateUpdateExtent();
    input->TriggerAsynchronousUpdate();
    input->UpdateData();
    }

  // Did the input change?
  // If it did then we should execute (i.e. we should send the data).
  // Note: We may need some logic to catch the case where the down port
  // has released its data.
  //if (downDataTime < input->GetMTime())
  if (input->GetDataReleased() == 0)
    {
    this->InvokeEvent(vtkCommand::StartEvent,NULL);
    // First transfer the new data.
    this->Controller->Send( input, remoteProcessId,
                            vtkInputPort::DATA_TRANSFER_TAG);
    this->InvokeEvent(vtkCommand::EndEvent,NULL);

    // Since this time has to be local to downstream process
    // and we have no data, we have to create a time here.
    // (The output data usually does this.)
    this->UpdateTime.Modified();

    // Since this OutputPort can have multiple InputPorts
    // and the InputPort makes the update-descision time comparison,
    // the InputPort has to store this time.
    downDataTime = this->UpdateTime.GetMTime();
    this->Controller->Send( &downDataTime, 1, remoteProcessId,
                            vtkInputPort::NEW_DATA_TIME_TAG);
    }
  else
    {  // Nothing to send.  We have to signal somehow.
    vtkDebugMacro("Promoting NULL (" << input << ") to process "
                  << remoteProcessId);
    this->Controller->Send( (vtkDataObject*)(NULL), remoteProcessId,
                            vtkInputPort::DATA_TRANSFER_TAG);

    // Go through the motions of sending the data time,
    // but just send the same data time back. (nothing changed).
    this->Controller->Send( &downDataTime, 1, remoteProcessId,
                            vtkInputPort::NEW_DATA_TIME_TAG);
    }

  // Postpone the update if we want pipeline parallism.
  // Handle no input gracefully. (Not true: Later we will send a NULL input.)
  if (this->PipelineFlag)
    {
    // change any parameters if the user wants to.
    if ( this->ParameterMethod )
      {
      (*this->ParameterMethod)(this->ParameterMethodArg);
      input->UpdateInformation();
      }

    // Update to anticipate the next request.
    if ( input != NULL )
      {
      input->UpdateInformation();
      input->PropagateUpdateExtent();
      input->TriggerAsynchronousUpdate();
      input->UpdateData();
      }
    }
}

//----------------------------------------------------------------------------
void vtkOutputPort::SetInput(vtkDataObject *input)
{
  if(input)
    {
    this->SetInputConnection(0, input->GetProducerPort());
    input->GetPipelineInformation()->Set(
      vtkDataObject::DATA_TYPE_NAME(), input->GetClassName());
    }
  else
    {
    // Setting a NULL input removes the connection.
    this->SetInputConnection(0, 0);
    }
}

//----------------------------------------------------------------------------
vtkDataObject *vtkOutputPort::GetInput()
{
  if (this->GetNumberOfInputConnections(0) < 1)
    {
    return 0;
    }
  return this->GetExecutive()->GetInputData(0, 0);
}

//----------------------------------------------------------------------------
// We need to create two RMIs when the tag is set.
// This means we must generate two tags form this ports tag.
// The ports tag should be even.
// (I do not like this, but is there another solution?)
void vtkOutputPort::SetTag(int tag)
{
  if (this->Tag == tag)
    {
    return;
    }

  this->Modified();

  // remove old RMI.
  if (this->Tag != -1)
    {
//    this->Controller->RemoveRMI(vtkOutputPortUpdateInformationCallBack,
//                                (void *)this, this->Tag);
//    this->Controller->RemoveRMI(vtkOutputPortUpdateCallBack,
//                                (void *)this, this->Tag + 1);
    }

  this->Tag = tag;
  this->Controller->AddRMI(vtkOutputPortUpdateInformationCallBack,
                           (void *)this, tag);
  this->Controller->AddRMI(vtkOutputPortUpdateCallBack,
                           (void *)this, tag+1);
}

//----------------------------------------------------------------------------
void vtkOutputPort::SetParameterMethod(void (*f)(void *), void *arg)
{
  if ( f != this->ParameterMethod || arg != this->ParameterMethodArg )
    {
    // delete the current arg if there is one and a delete meth
    if ((this->ParameterMethodArg)&&(this->ParameterMethodArgDelete))
      {
      (*this->ParameterMethodArgDelete)(this->ParameterMethodArg);
      }
    this->ParameterMethod = f;
    this->ParameterMethodArg = arg;
    this->Modified();
    }
}

//----------------------------------------------------------------------------
void vtkOutputPort::SetParameterMethodArgDelete(void (*f)(void *))
{
  if ( f != this->ParameterMethodArgDelete)
    {
    this->ParameterMethodArgDelete = f;
    this->Modified();
    }
}

//----------------------------------------------------------------------------
void vtkOutputPort::WaitForUpdate()
{
  this->Controller->ProcessRMIs();
}

//----------------------------------------------------------------------------
int vtkOutputPort
::FillInputPortInformation(int port, vtkInformation* info)
{
  info->Set(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkDataSet");
  return 1;
}
<commit_msg>COMP: fix compiler warning<commit_after>/*=========================================================================

  Program:   Visualization Toolkit
  Module:    vtkOutputPort.cxx

  Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
  All rights reserved.
  See Copyright.txt or http://www.kitware.com/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notice for more information.

=========================================================================*/
#include "vtkOutputPort.h"

#include "vtkCommand.h"
#include "vtkDataObject.h"
#include "vtkDemandDrivenPipeline.h"
#include "vtkInformation.h"
#include "vtkInputPort.h"
#include "vtkMultiProcessController.h"
#include "vtkMultiProcessController.h"
#include "vtkObjectFactory.h"

vtkCxxRevisionMacro(vtkOutputPort, "1.12");
vtkStandardNewMacro(vtkOutputPort);

vtkCxxSetObjectMacro(vtkOutputPort,Controller,vtkMultiProcessController);

//----------------------------------------------------------------------------
vtkOutputPort::vtkOutputPort()
{
  this->Tag = -1;

  // Controller keeps a reference to this object as well.
  this->Controller = NULL;
  this->SetController(vtkMultiProcessController::GetGlobalController());

  this->PipelineFlag = 0;
  this->ParameterMethod = NULL;
  this->ParameterMethodArgDelete = NULL;
  this->ParameterMethodArg = NULL;

  this->SetNumberOfInputPorts(1);
}

//----------------------------------------------------------------------------
// We need to have a "GetNetReferenceCount" to avoid memory leaks.
vtkOutputPort::~vtkOutputPort()
{
  this->SetController(0);

  if ((this->ParameterMethodArg)&&(this->ParameterMethodArgDelete))
    {
    (*this->ParameterMethodArgDelete)(this->ParameterMethodArg);
    }
}

//----------------------------------------------------------------------------
void vtkOutputPort::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os,indent);
  os << indent << "Tag: " << this->Tag << endl;
  os << indent << "Controller: (" << this->Controller << ")\n";
  os << indent << "Pipeline Flag: "
     << (this->PipelineFlag ? "On\n" : "Off\n");
}

//----------------------------------------------------------------------------
// Remote method call to UpdateInformation and send the information downstream.
// This should be a friend.
void vtkOutputPortUpdateInformationCallBack(void *arg, void *remoteArgs,
                                    int remoteArgsLength, int remoteProcessId)
{
  vtkOutputPort *self = (vtkOutputPort*)arg;

  remoteArgs = remoteArgs;
  remoteArgsLength = remoteArgsLength;
  // Just call a method
  self->TriggerUpdateInformation(remoteProcessId);
}
//----------------------------------------------------------------------------
void vtkOutputPort::TriggerUpdateInformation(int remoteProcessId)
{
  vtkDataObject *input = this->GetInput();

  // Handle no input gracefully.
  if ( input != NULL )
    {
    input->UpdateInformation();
    }

  // The MTime of the input should also be considered.
  // Important for pipeline parallelism.
  // Include it in the information for efficiency.
  unsigned long latestMTime;
  latestMTime = input->GetMTime();
  vtkDemandDrivenPipeline *ddp =
    vtkDemandDrivenPipeline::SafeDownCast(this->GetExecutive());
  if (ddp)
    {
    ddp->UpdateInformation();
    unsigned long t2 = ddp->GetPipelineMTime();
    if (t2 > latestMTime)
      {
      latestMTime = t2;
      }
    }

  // Now just send the information downstream.
  // PipelineMTime is part of information, so downstream
  // port will make the time comparison, and call Update if necessary.
  int wholeInformation[8];
  input->GetWholeExtent( wholeInformation );

  this->Controller->Send( wholeInformation, 7,
                          remoteProcessId, vtkInputPort::INFORMATION_TRANSFER_TAG);

  this->Controller->Send( &latestMTime, 1,
                          remoteProcessId, vtkInputPort::INFORMATION_TRANSFER_TAG );

  int maxNumPieces = input->GetMaximumNumberOfPieces();
  this->Controller->Send( &maxNumPieces, 1,
                          remoteProcessId, vtkInputPort::INFORMATION_TRANSFER_TAG );
}

//----------------------------------------------------------------------------
// Remote method call to Update and send data downstream.
// This should be a friend.
void vtkOutputPortUpdateCallBack(void *arg, void *remoteArgs,
                                 int remoteArgsLength, int remoteProcessId)
{
  vtkOutputPort *self = (vtkOutputPort*)arg;

  remoteArgs = remoteArgs;
  remoteArgsLength = remoteArgsLength;
  // Just call a method
  self->TriggerUpdate(remoteProcessId);
}
//----------------------------------------------------------------------------
void vtkOutputPort::TriggerUpdate(int remoteProcessId)
{
  unsigned long downDataTime;
  vtkDataObject *input = this->GetInput();

  // First get the update extent requested.
  int extent[9];
  this->Controller->Receive( extent, 9, remoteProcessId,
                            vtkInputPort::UPDATE_EXTENT_TAG);
  input->SetUpdateExtent( extent );
  input->SetUpdatePiece( extent[6] );
  input->SetUpdateNumberOfPieces( extent[7] );
  input->SetUpdateGhostLevel( extent[8] );

  // Note:  Receiving DataTime was the start of a more intelligent promotion
  // for pipeline parallism.  Unfortunately there was no way (I knew of)
  // for us to not promote on the first Update.  I backed off, and am
  // requiring either 1: Filters handle Update Not returning correct data,
  // or 2: Pipeline parallelism must be primed with Updates on the Output
  // ports (giving correct data).  This Receive can be removed.

  // This is for pipeline parallism.
  // This Output port may or may not promote our data (execute).
  // We need the data time of the last transfer to compare to the mtime
  // of our input to determine if it should send the data (execute).
  this->Controller->Receive( &(downDataTime), 1, remoteProcessId,
                             vtkInputPort::NEW_DATA_TIME_TAG);

  // What was the idea of using relesed data here?  It caused a bug for multiple updates.
  // if ( input != NULL && input->GetDataReleased())

  // Postpone the update if we want pipeline parallism.
  // Handle no input gracefully. (Not true: Later we will send a NULL input.)
  if ( input != NULL && this->PipelineFlag == 0)
    {
    input->UpdateInformation();
    input->PropagateUpdateExtent();
    input->TriggerAsynchronousUpdate();
    input->UpdateData();
    }

  // Did the input change?
  // If it did then we should execute (i.e. we should send the data).
  // Note: We may need some logic to catch the case where the down port
  // has released its data.
  //if (downDataTime < input->GetMTime())
  if (input->GetDataReleased() == 0)
    {
    this->InvokeEvent(vtkCommand::StartEvent,NULL);
    // First transfer the new data.
    this->Controller->Send( input, remoteProcessId,
                            vtkInputPort::DATA_TRANSFER_TAG);
    this->InvokeEvent(vtkCommand::EndEvent,NULL);

    // Since this time has to be local to downstream process
    // and we have no data, we have to create a time here.
    // (The output data usually does this.)
    this->UpdateTime.Modified();

    // Since this OutputPort can have multiple InputPorts
    // and the InputPort makes the update-descision time comparison,
    // the InputPort has to store this time.
    downDataTime = this->UpdateTime.GetMTime();
    this->Controller->Send( &downDataTime, 1, remoteProcessId,
                            vtkInputPort::NEW_DATA_TIME_TAG);
    }
  else
    {  // Nothing to send.  We have to signal somehow.
    vtkDebugMacro("Promoting NULL (" << input << ") to process "
                  << remoteProcessId);
    this->Controller->Send( (vtkDataObject*)(NULL), remoteProcessId,
                            vtkInputPort::DATA_TRANSFER_TAG);

    // Go through the motions of sending the data time,
    // but just send the same data time back. (nothing changed).
    this->Controller->Send( &downDataTime, 1, remoteProcessId,
                            vtkInputPort::NEW_DATA_TIME_TAG);
    }

  // Postpone the update if we want pipeline parallism.
  // Handle no input gracefully. (Not true: Later we will send a NULL input.)
  if (this->PipelineFlag)
    {
    // change any parameters if the user wants to.
    if ( this->ParameterMethod )
      {
      (*this->ParameterMethod)(this->ParameterMethodArg);
      input->UpdateInformation();
      }

    // Update to anticipate the next request.
    if ( input != NULL )
      {
      input->UpdateInformation();
      input->PropagateUpdateExtent();
      input->TriggerAsynchronousUpdate();
      input->UpdateData();
      }
    }
}

//----------------------------------------------------------------------------
void vtkOutputPort::SetInput(vtkDataObject *input)
{
  if(input)
    {
    this->SetInputConnection(0, input->GetProducerPort());
    input->GetPipelineInformation()->Set(
      vtkDataObject::DATA_TYPE_NAME(), input->GetClassName());
    }
  else
    {
    // Setting a NULL input removes the connection.
    this->SetInputConnection(0, 0);
    }
}

//----------------------------------------------------------------------------
vtkDataObject *vtkOutputPort::GetInput()
{
  if (this->GetNumberOfInputConnections(0) < 1)
    {
    return 0;
    }
  return this->GetExecutive()->GetInputData(0, 0);
}

//----------------------------------------------------------------------------
// We need to create two RMIs when the tag is set.
// This means we must generate two tags form this ports tag.
// The ports tag should be even.
// (I do not like this, but is there another solution?)
void vtkOutputPort::SetTag(int tag)
{
  if (this->Tag == tag)
    {
    return;
    }

  this->Modified();

  // remove old RMI.
  if (this->Tag != -1)
    {
//    this->Controller->RemoveRMI(vtkOutputPortUpdateInformationCallBack,
//                                (void *)this, this->Tag);
//    this->Controller->RemoveRMI(vtkOutputPortUpdateCallBack,
//                                (void *)this, this->Tag + 1);
    }

  this->Tag = tag;
  this->Controller->AddRMI(vtkOutputPortUpdateInformationCallBack,
                           (void *)this, tag);
  this->Controller->AddRMI(vtkOutputPortUpdateCallBack,
                           (void *)this, tag+1);
}

//----------------------------------------------------------------------------
void vtkOutputPort::SetParameterMethod(void (*f)(void *), void *arg)
{
  if ( f != this->ParameterMethod || arg != this->ParameterMethodArg )
    {
    // delete the current arg if there is one and a delete meth
    if ((this->ParameterMethodArg)&&(this->ParameterMethodArgDelete))
      {
      (*this->ParameterMethodArgDelete)(this->ParameterMethodArg);
      }
    this->ParameterMethod = f;
    this->ParameterMethodArg = arg;
    this->Modified();
    }
}

//----------------------------------------------------------------------------
void vtkOutputPort::SetParameterMethodArgDelete(void (*f)(void *))
{
  if ( f != this->ParameterMethodArgDelete)
    {
    this->ParameterMethodArgDelete = f;
    this->Modified();
    }
}

//----------------------------------------------------------------------------
void vtkOutputPort::WaitForUpdate()
{
  this->Controller->ProcessRMIs();
}

//----------------------------------------------------------------------------
int vtkOutputPort
::FillInputPortInformation(int vtkNotUsed( port ), vtkInformation* info)
{
  info->Set(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkDataSet");
  return 1;
}
<|endoftext|>


<commit_before>//  this is so we can implment hash function for dynamic_bitset
#define BOOST_DYNAMIC_BITSET_DONT_USE_FRIENDS

#include <iostream>
#include <string>
#include <boost/program_options.hpp>
#include <vector>
#include <fstream>
#include "ioHandler.h"
#include <map>
#include <unordered_map>
#include <boost/dynamic_bitset.hpp>
#include <boost/iostreams/filter/gzip.hpp>
#include <boost/iostreams/filtering_stream.hpp>
#include <boost/iostreams/device/file_descriptor.hpp>
#include <boost/iostreams/stream.hpp>
#include <boost/iostreams/device/file.hpp>
#include <boost/iostreams/device/mapped_file.hpp>
#include <boost/filesystem/path.hpp>
#include <boost/filesystem.hpp>
#include <boost/functional/hash.hpp>

namespace
{
    const size_t SUCCESS = 0;
    const size_t ERROR_IN_COMMAND_LINE = 1;
    const size_t ERROR_UNHANDLED_EXCEPTION = 2;

} // namespace

class dbhash {
public:
    std::size_t operator() (const boost::dynamic_bitset<>& bs) const {
        return boost::hash_value(bs.m_bits);
    }
};

typedef std::unordered_map<std::string, size_t> Counter;
typedef std::unordered_map <boost::dynamic_bitset<>, std::unique_ptr<ReadBase>, dbhash> BitMap;

template <class T, class Impl>
void load_map(InputReader<T, Impl> &reader, Counter& counters, BitMap& read_map, size_t start, size_t length) {
    while(reader.has_next()) {
        auto i = reader.next();
        ++counters["TotalRecords"];
        //check for existance, store or compare quality and replace:
        if (auto key=i->get_key(start, length)) {
            // find faster than count on some compilers
            if(read_map.find(*key) == read_map.end()) {
                read_map[*key] = std::move(i);
            } else if(i->avg_q_score() > read_map[*key]->avg_q_score()){
                read_map[*key] = std::move(i);
                ++counters["Replaced"];
            }
        } else {  // key had N
            ++counters["HasN"];
        }
    }
}

template <class T>
void output_read_map_fastq(const BitMap& read_map, T& out1, T& out2, T& single) {

    OutputWriter<PairedEndRead, PairedEndReadOutFastq> pofs(out1, out2);
    OutputWriter<SingleEndRead, SingleEndReadOutFastq> sofs(single);
    for(auto const &i : read_map) {
        PairedEndRead* per = dynamic_cast<PairedEndRead*>(i.second.get());
        if (per) {
            pofs.write(*per);
        } else {
            SingleEndRead* ser = dynamic_cast<SingleEndRead*>(i.second.get());
            if(ser) {
                sofs.write(*ser);
            }
            else {
                throw std::runtime_error("Unkown read found");
            }
        }
    }
}

namespace bi = boost::iostreams;
namespace bf = boost::filesystem;

int check_open_r(const std::string& filename) {
    bf::path p(filename);
    if (!bf::exists(p)) {
        throw std::runtime_error("File " + filename + " was not found.");
    }

    if (p.extension() == ".gz") {
        return fileno(popen(("gunzip -c " + filename).c_str(), "r"));
    } else {
        return fileno(fopen(filename.c_str(), "r"));
    }
}

int main(int argc, char** argv)
{
    BitMap read_map;
    Counter counters;
    counters["TotalRecords"] = 0;
    counters["Replaced"] = 0;
    counters["HasN"] = 0;
    size_t start, length = 0;
    std::string prefix;
    std::vector<std::string> default_outfiles = {"PE1", "PE2", "SE"};
    bool fastq_out;
    bool tab_out;
    bool std_out;
    bool gzip_out;

    try
    {
        /** Define and parse the program options
         */
        namespace po = boost::program_options;
        po::options_description desc("Options");
        desc.add_options()
            ("version,v",                  "Version print")
            ("read1-input,1", po::value< std::vector<std::string> >(),
                                           "Read 1 input <comma sep for multiple files>")
            ("read2-input,2", po::value< std::vector<std::string> >(),
                                           "Read 2 input <comma sep for multiple files>")
            ("singleend-input,U", po::value< std::vector<std::string> >(),
                                           "Single end read input <comma sep for multiple files>")
            ("tab-input,T", po::value< std::vector<std::string> >(),
                                           "Tab input <comma sep for multiple files>")
            ("interleaved-input,I", po::value< std::vector<std::string> >(),
                                           "Interleaved input I <comma sep for multiple files>")
            ("stdin-input,S",              "STDIN input <MUST BE TAB DELIMITED INPUT>")
            ("start,s", po::value<size_t>(&start)->default_value(10),  "Start location for unique ID <int>")
            ("length,l", po::value<size_t>(&length)->default_value(10), "Length of unique ID <int>")
            ("quality-check-off,q",        "Quality Checking Off First Duplicate seen will be kept")
            ("gzip-output,g", po::bool_switch(&gzip_out)->default_value(false),  "Output gzipped")
            ("interleaved-output, i",      "Output to interleaved")
            ("fastq-output,f", po::bool_switch(&fastq_out)->default_value(false), "Fastq format output")
            ("force,F", po::bool_switch()->default_value(true),         "Forces overwrite of files")
            ("tab-output,t", po::bool_switch(&tab_out)->default_value(false),   "Tab-delimited output")
            ("to-stdout,O", po::bool_switch(&std_out)->default_value(false),    "Prints to STDOUT in Tab Delimited")
            ("prefix,p", po::value<std::string>(&prefix)->default_value("output_nodup_"),
                                           "Prefix for outputted files")
            ("log-file,L",                 "Output-Logfile")
            ("no-log,N",                   "No logfile <outputs to stderr>")
            ("help,h",                     "Prints help.");

        po::variables_map vm;
        try
        {
            po::store(po::parse_command_line(argc, argv, desc),
                      vm); // can throw

            /** --help option
             */
            if ( vm.count("help")  || vm.size() == 0)
            {
                std::cout << "Super-Deduper" << std::endl
                          << desc << std::endl;
                return SUCCESS;
            }

            po::notify(vm); // throws on error, so do after help in case
            // there are any problems
            if(vm.count("read1-input")) {
                if (!vm.count("read2-input")) {
                    throw std::runtime_error("must specify both read1 and read2 input files.");
                } else if (vm.count("read2-input") != vm.count("read1-input")) {
                    throw std::runtime_error("must have same number of input files for read1 and read2");
                }
                auto read1_files = vm["read1-input"].as<std::vector<std::string> >();
                auto read2_files = vm["read2-input"].as<std::vector<std::string> >();

                for(size_t i = 0; i < read1_files.size(); ++i) {
                    bi::stream<bi::file_descriptor_source> is1{check_open_r(read1_files[i]), bi::close_handle};
                    bi::stream<bi::file_descriptor_source> is2{check_open_r(read2_files[i]), bi::close_handle};

                    InputReader<PairedEndRead, PairedEndReadFastqImpl> ifp(is1, is2);
                    load_map(ifp, counters, read_map, start, length);
                }
            }
            if(vm.count("singleend-input")) {
                auto read_files = vm["singleend-input"].as<std::vector<std::string> >();
                for (auto file : read_files) {
                    std::ifstream read1(file, std::ifstream::in);
                    InputReader<SingleEndRead, SingleEndReadFastqImpl> ifs(read1);
                    load_map(ifs, counters, read_map, start, length);
                }
            }

            if(vm.count("tab-input")) {
                auto read_files = vm["tab-input"].as<std::vector<std::string> > ();
                for (auto file : read_files) {
                    std::ifstream read1(file, std::ifstream::in);
                }
            }

            if (fastq_out || (! std_out && ! tab_out) ) {
                for (auto& outfile: default_outfiles) {
                    outfile = prefix + outfile + ".fastq";
                }

                if (gzip_out) {
                    bi::stream<bi::file_descriptor_sink> out1{fileno(popen(("gzip > " + default_outfiles[0] + ".gz").c_str(), "w")), bi::close_handle};
                    bi::stream<bi::file_descriptor_sink> out2{fileno(popen(("gzip > " + default_outfiles[1] + ".gz").c_str(), "w")), bi::close_handle};
                    bi::stream<bi::file_descriptor_sink> out3{fileno(popen(("gzip > " + default_outfiles[2] + ".gz").c_str(), "w")), bi::close_handle};

                    output_read_map_fastq(read_map, out1, out2, out3);

                } else {
                    // note: mapped file is faster but uses more memory
                    std::ofstream out1(default_outfiles[0], std::ofstream::out);
                    std::ofstream out2(default_outfiles[1], std::ofstream::out);
                    std::ofstream out3(default_outfiles[2], std::ofstream::out);
                    //bi::stream<bi::mapped_file_sink> out1{default_outfiles[0].c_str()};
                    //bi::stream<bi::mapped_file_sink> out2{default_outfiles[1].c_str()};
                    //bi::stream<bi::mapped_file_sink> out3{default_outfiles[2].c_str()};
                    output_read_map_fastq(read_map, out1, out2, out3);
                }
            }

        }
        catch(po::error& e)
        {
            std::cerr << "ERROR: " << e.what() << std::endl << std::endl;
            std::cerr << desc << std::endl;
            return ERROR_IN_COMMAND_LINE;
        }

    }
    catch(std::exception& e)
    {
        std::cerr << "\n\tUnhandled Exception: "
                  << e.what() << std::endl;
        return ERROR_UNHANDLED_EXCEPTION;

    }

    std::cerr << "TotalRecords:" << counters["TotalRecords"] << "\tReplaced:" << counters["Replaced"]
              << "\tKept:" << read_map.size() << "\tRemoved:" << counters["TotalRecords"] - read_map.size()
              << "\tHasN:" << counters["HasN"] << std::endl;
    return SUCCESS;

}
<commit_msg>Allow SD to read and write tab and inter files<commit_after>//  this is so we can implment hash function for dynamic_bitset
#define BOOST_DYNAMIC_BITSET_DONT_USE_FRIENDS

#include <iostream>
#include <string>
#include <boost/program_options.hpp>
#include <vector>
#include <fstream>
#include "ioHandler.h"
#include <map>
#include <unordered_map>
#include <boost/dynamic_bitset.hpp>
#include <boost/iostreams/filter/gzip.hpp>
#include <boost/iostreams/filtering_stream.hpp>
#include <boost/iostreams/device/file_descriptor.hpp>
#include <boost/iostreams/stream.hpp>
#include <boost/iostreams/device/file.hpp>
#include <boost/iostreams/device/mapped_file.hpp>
#include <boost/filesystem/path.hpp>
#include <boost/filesystem.hpp>
#include <boost/functional/hash.hpp>

namespace
{
    const size_t SUCCESS = 0;
    const size_t ERROR_IN_COMMAND_LINE = 1;
    const size_t ERROR_UNHANDLED_EXCEPTION = 2;

} // namespace

class dbhash {
public:
    std::size_t operator() (const boost::dynamic_bitset<>& bs) const {
        return boost::hash_value(bs.m_bits);
    }
};

typedef std::unordered_map<std::string, size_t> Counter;
typedef std::unordered_map <boost::dynamic_bitset<>, std::unique_ptr<ReadBase>, dbhash> BitMap;

template <class T, class Impl>
void load_map(InputReader<T, Impl> &reader, Counter& counters, BitMap& read_map, size_t start, size_t length) {
    while(reader.has_next()) {
        auto i = reader.next();
        ++counters["TotalRecords"];
        //check for existance, store or compare quality and replace:
        if (auto key=i->get_key(start, length)) {
            // find faster than count on some compilers
            if(read_map.find(*key) == read_map.end()) {
                read_map[*key] = std::move(i);
            } else if(i->avg_q_score() > read_map[*key]->avg_q_score()){
                read_map[*key] = std::move(i);
                ++counters["Replaced"];
            }
        } else {  // key had N
            ++counters["HasN"];
        }
    }
}

template <class T>
void output_read_map_tab(const BitMap& read_map, T& out1) {
    OutputWriter<ReadBase, ReadBaseOutTab> tabs(out1);
    for(auto const &i : read_map) {
        ReadBase* rb = dynamic_cast<ReadBase*>(i.second.get());
        tabs.write(*rb);
    }
}

template <class T>
void output_read_map_inter(const BitMap& read_map, T& out1, T& single) {
    OutputWriter<PairedEndRead, PairedEndReadOutInter> pofs(out1);
    OutputWriter<SingleEndRead, SingleEndReadOutFastq> sofs(single);
    for(auto const &i : read_map) {
        PairedEndRead* per = dynamic_cast<PairedEndRead*>(i.second.get());
        if (per) {
            pofs.write(*per);
        } else {
            SingleEndRead* ser = dynamic_cast<SingleEndRead*>(i.second.get());
            if(ser) {
                sofs.write(*ser);
            }
            else {
                throw std::runtime_error("Unkown read found");
            }
        }
    }
}

template <class T>
void output_read_map_fastq(const BitMap& read_map, T& out1, T& out2, T& single) {

    OutputWriter<PairedEndRead, PairedEndReadOutFastq> pofs(out1, out2);
    OutputWriter<SingleEndRead, SingleEndReadOutFastq> sofs(single);
    for(auto const &i : read_map) {
        PairedEndRead* per = dynamic_cast<PairedEndRead*>(i.second.get());
        if (per) {
            pofs.write(*per);
        } else {
            SingleEndRead* ser = dynamic_cast<SingleEndRead*>(i.second.get());
            if(ser) {
                sofs.write(*ser);
            }
            else {
                throw std::runtime_error("Unkown read found");
            }
        }
    }
}

namespace bi = boost::iostreams;
namespace bf = boost::filesystem;

int check_open_r(const std::string& filename) {
    bf::path p(filename);
    if (!bf::exists(p)) {
        throw std::runtime_error("File " + filename + " was not found.");
    }

    if (p.extension() == ".gz") {
        return fileno(popen(("gunzip -c " + filename).c_str(), "r"));
    } else {
        return fileno(fopen(filename.c_str(), "r"));
    }
}

int main(int argc, char** argv)
{
    BitMap read_map;
    Counter counters;
    counters["TotalRecords"] = 0;
    counters["Replaced"] = 0;
    counters["HasN"] = 0;
    size_t start, length = 0;
    std::string prefix;
    std::vector<std::string> default_outfiles = {"PE1", "PE2", "SE"};
    bool fastq_out;
    bool tab_out;
    bool std_out;
    bool gzip_out;
    bool interleaved_out;

    try
    {
        /** Define and parse the program options
         */
        namespace po = boost::program_options;
        po::options_description desc("Options");
        desc.add_options()
            ("version,v",                  "Version print")
            ("read1-input,1", po::value< std::vector<std::string> >(),
                                           "Read 1 input <comma sep for multiple files>")
            ("read2-input,2", po::value< std::vector<std::string> >(),
                                           "Read 2 input <comma sep for multiple files>")
            ("singleend-input,U", po::value< std::vector<std::string> >(),
                                           "Single end read input <comma sep for multiple files>")
            ("tab-input,T", po::value< std::vector<std::string> >(),
                                           "Tab input <comma sep for multiple files>")
            ("interleaved-input,I", po::value< std::vector<std::string> >(),
                                           "Interleaved input I <comma sep for multiple files>")
            ("stdin-input,S",              "STDIN input <MUST BE TAB DELIMITED INPUT>")
            ("start,s", po::value<size_t>(&start)->default_value(10),  "Start location for unique ID <int>")
            ("length,l", po::value<size_t>(&length)->default_value(10), "Length of unique ID <int>")
            ("quality-check-off,q",        "Quality Checking Off First Duplicate seen will be kept")
            ("gzip-output,g", po::bool_switch(&gzip_out)->default_value(false),  "Output gzipped")
            ("interleaved-output, i", po::bool_switch(&interleaved_out)->default_value(false),     "Output to interleaved")
            ("fastq-output,f", po::bool_switch(&fastq_out)->default_value(false), "Fastq format output")
            ("force,F", po::bool_switch()->default_value(true),         "Forces overwrite of files")
            ("tab-output,t", po::bool_switch(&tab_out)->default_value(false),   "Tab-delimited output")
            ("to-stdout,O", po::bool_switch(&std_out)->default_value(false),    "Prints to STDOUT in Tab Delimited")
            ("prefix,p", po::value<std::string>(&prefix)->default_value("output_nodup_"),
                                           "Prefix for outputted files")
            ("log-file,L",                 "Output-Logfile")
            ("no-log,N",                   "No logfile <outputs to stderr>")
            ("help,h",                     "Prints help.");

        po::variables_map vm;
        try
        {
            po::store(po::parse_command_line(argc, argv, desc),
                      vm); // can throw

            /** --help option
             */
            if ( vm.count("help")  || vm.size() == 0)
            {
                std::cout << "Super-Deduper" << std::endl
                          << desc << std::endl;
                return SUCCESS;
            }

            po::notify(vm); // throws on error, so do after help in case
            // there are any problems
            if(vm.count("read1-input")) {
                if (!vm.count("read2-input")) {
                    throw std::runtime_error("must specify both read1 and read2 input files.");
                } else if (vm.count("read2-input") != vm.count("read1-input")) {
                    throw std::runtime_error("must have same number of input files for read1 and read2");
                }
                auto read1_files = vm["read1-input"].as<std::vector<std::string> >();
                auto read2_files = vm["read2-input"].as<std::vector<std::string> >();

                for(size_t i = 0; i < read1_files.size(); ++i) {
                    bi::stream<bi::file_descriptor_source> is1{check_open_r(read1_files[i]), bi::close_handle};
                    bi::stream<bi::file_descriptor_source> is2{check_open_r(read2_files[i]), bi::close_handle};

                    InputReader<PairedEndRead, PairedEndReadFastqImpl> ifp(is1, is2);
                    load_map(ifp, counters, read_map, start, length);
                }
            }
            if(vm.count("singleend-input")) {
                auto read_files = vm["singleend-input"].as<std::vector<std::string> >();
                for (auto file : read_files) {
                    std::ifstream read1(file, std::ifstream::in);
                    InputReader<SingleEndRead, SingleEndReadFastqImpl> ifs(read1);
                    load_map(ifs, counters, read_map, start, length);
                }
            }

            if(vm.count("tab-input")) {
                auto read_files = vm["tab-input"].as<std::vector<std::string> > ();
                for (auto file : read_files) {
                    std::ifstream tabReads(file, std::ifstream::in);
                    InputReader<ReadBase, TabReadImpl> ift(tabReads);
                    load_map(ift, counters, read_map, start, length);
                }
            }

            if (vm.count("interleaved-input")) {
                auto read_files = vm["interleaved-input"].as<std::vector<std::string > >();
                for (auto file : read_files) {
                    std::ifstream interReads(file, std::ifstream::in);
                    InputReader<PairedEndRead, InterReadImpl> ifp(interReads);
                    load_map(ifp, counters, read_map, start, length);
                }
            }

            if (fastq_out || (! std_out && ! tab_out) ) {
                for (auto& outfile: default_outfiles) {
                    outfile = prefix + outfile + ".fastq";
                }

                if (gzip_out) {
                    bi::stream<bi::file_descriptor_sink> out1{fileno(popen(("gzip > " + default_outfiles[0] + ".gz").c_str(), "w")), bi::close_handle};
                    bi::stream<bi::file_descriptor_sink> out2{fileno(popen(("gzip > " + default_outfiles[1] + ".gz").c_str(), "w")), bi::close_handle};
                    bi::stream<bi::file_descriptor_sink> out3{fileno(popen(("gzip > " + default_outfiles[2] + ".gz").c_str(), "w")), bi::close_handle};
                    output_read_map_fastq(read_map, out1, out2, out3);
                } else {
                    // note: mapped file is faster but uses more memory
                    std::ofstream out1(default_outfiles[0], std::ofstream::out);
                    std::ofstream out2(default_outfiles[1], std::ofstream::out);
                    std::ofstream out3(default_outfiles[2], std::ofstream::out);
                    //bi::stream<bi::mapped_file_sink> out1{default_outfiles[0].c_str()};
                    //bi::stream<bi::mapped_file_sink> out2{default_outfiles[1].c_str()};
                    //bi::stream<bi::mapped_file_sink> out3{default_outfiles[2].c_str()};
                    output_read_map_fastq(read_map, out1, out2, out3);
                }
            } else if (interleaved_out)  {
                for (auto& outfile: default_outfiles) {
                    outfile = prefix + "INTER" + ".fastq";
                }

                if (gzip_out) {
                    bi::stream<bi::file_descriptor_sink> out1{fileno(popen(("gzip > " + default_outfiles[0] + ".gz").c_str(), "w")), bi::close_handle};
                    bi::stream<bi::file_descriptor_sink> out3{fileno(popen(("gzip > " + default_outfiles[2] + ".gz").c_str(), "w")), bi::close_handle};
                    output_read_map_inter(read_map, out1, out3);
                } else {
                    std::ofstream out1(default_outfiles[0], std::ofstream::out);
                    std::ofstream out3(default_outfiles[2], std::ofstream::out);
                    output_read_map_inter(read_map, out1, out3);
                }
            } else if (tab_out) {
                for (auto& outfile: default_outfiles) {
                    outfile = prefix + "tab" + ".tastq";
                }

                if (gzip_out) {
                    bi::stream<bi::file_descriptor_sink> out1{fileno(popen(("gzip > " + default_outfiles[0] + ".gz").c_str(), "w")), bi::close_handle};
                    output_read_map_tab(read_map, out1);
                } else {
                    std::ofstream out1(default_outfiles[0], std::ofstream::out);
                    output_read_map_tab(read_map, out1);
                }
            }

        }
        catch(po::error& e)
        {
            std::cerr << "ERROR: " << e.what() << std::endl << std::endl;
            std::cerr << desc << std::endl;
            return ERROR_IN_COMMAND_LINE;
        }

    }
    catch(std::exception& e)
    {
        std::cerr << "\n\tUnhandled Exception: "
                  << e.what() << std::endl;
        return ERROR_UNHANDLED_EXCEPTION;

    }

    std::cerr << "TotalRecords:" << counters["TotalRecords"] << "\tReplaced:" << counters["Replaced"]
              << "\tKept:" << read_map.size() << "\tRemoved:" << counters["TotalRecords"] - read_map.size()
              << "\tHasN:" << counters["HasN"] << std::endl;
    return SUCCESS;

}
<|endoftext|>


<commit_before>// Copyright (c) 2011 Zeex
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "amxprofiler.h"

platformstl::int64_t AMXFunPerfCounter::GetExecutionTime() const {
    return time_;
}

platformstl::int64_t AMXFunPerfCounter::GetNumberOfCalls() const {
    return calls_;
}

void AMXFunPerfCounter::StartCounter() {
    counter_.start();
}

void AMXFunPerfCounter::StopCounter() {
    counter_.stop();
    time_ += counter_.get_microseconds();
}

void AMXFunPerfCounter::IncreaseCalls() {
    calls_++;
}

AMXProfiler::AMXProfiler(AMX *amx)
    : amx_(amx),
      running_(false),
      debugHook_(amx->debug),
      callback_(amx->callback),
      currentStackFrame_(0)
{
    amx_SetUserData(amx_, AMX_USERTAG('p', 'r', 'o', 'f'), this);
}

AMXProfiler::~AMXProfiler() {
    amx_SetUserData(amx_, AMX_USERTAG('p', 'r', 'o', 'f'), 0);
}

static int AMXAPI DebugHook(AMX *amx) {
    void *prof;
    amx_GetUserData(amx, AMX_USERTAG('p', 'r', 'o', 'f'), &prof);
    return static_cast<AMXProfiler*>(prof)->DebugHook();
}

static int AMXAPI Callback(AMX *amx, cell index, cell *result, cell *params) {
    void *prof;
    amx_GetUserData(amx, AMX_USERTAG('p', 'r', 'o', 'f'), &prof);
    return static_cast<AMXProfiler*>(prof)->Callback(index, result, params);
}

void AMXProfiler::Run() {
    running_ = true;
    functions_.clear();
    amx_SetDebugHook(amx_, ::DebugHook);
    amx_SetCallback(amx_, ::Callback);
}

bool AMXProfiler::IsRunning() const {
    return running_;
}

void AMXProfiler::Terminate() {
    running_ = false;
    amx_SetDebugHook(amx_, debugHook_);
    amx_SetCallback(amx_, callback_);
}

// Get address of a CALL by frame
static cell GetCallAddress(AMX *amx, cell frame) {
    AMX_HEADER *hdr = reinterpret_cast<AMX_HEADER*>(amx->base);
    cell data = reinterpret_cast<cell>(amx->base + hdr->dat);
    return *(reinterpret_cast<cell*>(data + frame) + 1) - 2*sizeof(cell);
}

// Get address of callee
static cell GetCallTarget(AMX *amx, cell callAddr) {
    AMX_HEADER *hdr = reinterpret_cast<AMX_HEADER*>(amx->base);
    cell code = reinterpret_cast<cell>(amx->base) + hdr->cod;
    return *reinterpret_cast<cell*>(callAddr + sizeof(cell) + code) - code;
}

int AMXProfiler::DebugHook() {
    if (amx_->frm != currentStackFrame_) {
        if (amx_->frm < currentStackFrame_ || currentStackFrame_ == 0) {
            // Entered a function
            cell callAddr = GetCallAddress(amx_, amx_->frm);
            if (callAddr > 0) {
                cell funcAddr = GetCallTarget(amx_, callAddr);
                functions_[funcAddr].IncreaseCalls();
                functions_[funcAddr].StartCounter();
            }
        } else if (amx_->frm > currentStackFrame_) {
            // Left a function
            cell callAddr = GetCallAddress(amx_, currentStackFrame_);
            if (callAddr > 0) {
                cell funcAddr = GetCallTarget(amx_, callAddr);
                functions_[funcAddr].StopCounter();
            }
        }
        currentStackFrame_ = amx_->frm;
    }

    if (debugHook_ != 0) {
        // Others could set their own debug hooks
        return debugHook_(amx_);
    }
    return AMX_ERR_NONE;
}

int AMXProfiler::Callback(cell index, cell *result, cell *params) {
    AMXFunPerfCounter &fun = functions_[-index]; // Note negative index
    fun.StartCounter();

    // Disable SYSREQ.D opcodes
    amx_->sysreq_d = 0;

    // Call any previously set AMX callback (must not be null
    // so we don't bother checking it)
    int error = callback_(amx_, index, result, params);

    fun.StopCounter();
    fun.IncreaseCalls();

    return error;
}

std::vector<AMXFunPerfStats> AMXProfiler::GetStats() const {
    std::vector<AMXFunPerfStats> stats;

    for (std::map<cell, AMXFunPerfCounter>::const_iterator it = functions_.begin();
         it != functions_.end(); ++it)
    {
        AMXFunPerfStats st;
        cell address = it->first;
        st.native = address <= 0;
        st.address = address < 0 ? -address : address;
        st.numberOfCalls = it->second.GetNumberOfCalls();
        st.executionTime = it->second.GetExecutionTime();
        stats.push_back(st);
    }

    return stats;
}

<commit_msg>Fix sometimes Overall Time % is shown as 0.00 but it's not (numberOfCalls == 0)<commit_after>// Copyright (c) 2011 Zeex
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "amxprofiler.h"

platformstl::int64_t AMXFunPerfCounter::GetExecutionTime() const {
    return time_;
}

platformstl::int64_t AMXFunPerfCounter::GetNumberOfCalls() const {
    return calls_;
}

void AMXFunPerfCounter::StartCounter() {
    counter_.start();
}

void AMXFunPerfCounter::StopCounter() {
    counter_.stop();
    time_ += counter_.get_microseconds();
}

void AMXFunPerfCounter::IncreaseCalls() {
    calls_++;
}

AMXProfiler::AMXProfiler(AMX *amx)
    : amx_(amx),
      running_(false),
      debugHook_(amx->debug),
      callback_(amx->callback),
      currentStackFrame_(0)
{
    amx_SetUserData(amx_, AMX_USERTAG('p', 'r', 'o', 'f'), this);
}

AMXProfiler::~AMXProfiler() {
    amx_SetUserData(amx_, AMX_USERTAG('p', 'r', 'o', 'f'), 0);
}

static int AMXAPI DebugHook(AMX *amx) {
    void *prof;
    amx_GetUserData(amx, AMX_USERTAG('p', 'r', 'o', 'f'), &prof);
    return static_cast<AMXProfiler*>(prof)->DebugHook();
}

static int AMXAPI Callback(AMX *amx, cell index, cell *result, cell *params) {
    void *prof;
    amx_GetUserData(amx, AMX_USERTAG('p', 'r', 'o', 'f'), &prof);
    return static_cast<AMXProfiler*>(prof)->Callback(index, result, params);
}

void AMXProfiler::Run() {
    running_ = true;
    functions_.clear();
    amx_SetDebugHook(amx_, ::DebugHook);
    amx_SetCallback(amx_, ::Callback);
}

bool AMXProfiler::IsRunning() const {
    return running_;
}

void AMXProfiler::Terminate() {
    running_ = false;
    amx_SetDebugHook(amx_, debugHook_);
    amx_SetCallback(amx_, callback_);
}

// Get address of a CALL by frame
static cell GetCallAddress(AMX *amx, cell frame) {
    AMX_HEADER *hdr = reinterpret_cast<AMX_HEADER*>(amx->base);
    cell data = reinterpret_cast<cell>(amx->base + hdr->dat);
    return *(reinterpret_cast<cell*>(data + frame) + 1) - 2*sizeof(cell);
}

// Get address of callee
static cell GetCallTarget(AMX *amx, cell callAddr) {
    AMX_HEADER *hdr = reinterpret_cast<AMX_HEADER*>(amx->base);
    cell code = reinterpret_cast<cell>(amx->base) + hdr->cod;
    return *reinterpret_cast<cell*>(callAddr + sizeof(cell) + code) - code;
}

int AMXProfiler::DebugHook() {
    if (amx_->frm != currentStackFrame_) {
        if (amx_->frm < currentStackFrame_ || currentStackFrame_ == 0) {
            // Entered a function
            cell callAddr = GetCallAddress(amx_, amx_->frm);
            if (callAddr > 0) {
                cell funcAddr = GetCallTarget(amx_, callAddr);
                functions_[funcAddr].IncreaseCalls();
                functions_[funcAddr].StartCounter();
            }
        } else if (amx_->frm > currentStackFrame_) {
            // Left a function
            cell callAddr = GetCallAddress(amx_, currentStackFrame_);
            if (callAddr > 0) {
                cell funcAddr = GetCallTarget(amx_, callAddr);
                functions_[funcAddr].StopCounter();
            }
        }
        currentStackFrame_ = amx_->frm;
    }

    if (debugHook_ != 0) {
        // Others could set their own debug hooks
        return debugHook_(amx_);
    }
    return AMX_ERR_NONE;
}

int AMXProfiler::Callback(cell index, cell *result, cell *params) {
    AMXFunPerfCounter &fun = functions_[-index]; // Notice negative index
    fun.StartCounter();

    // The default AMX callback (amx_Callback) can replace SYSREQ.C opcodes
    // with SYSREQ.D for better performance.
    amx_->sysreq_d = 0;

    // Call any previously set AMX callback (must not be null so we don't check)
    int error = callback_(amx_, index, result, params);

    fun.StopCounter();
    fun.IncreaseCalls();

    return error;
}

std::vector<AMXFunPerfStats> AMXProfiler::GetStats() const {
    std::vector<AMXFunPerfStats> stats;

    for (std::map<cell, AMXFunPerfCounter>::const_iterator it = functions_.begin();
         it != functions_.end(); ++it)
    {
        if (it->second.GetNumberOfCalls() > 0) { // Sometimes calls == 0
            AMXFunPerfStats st;
            cell address = it->first;
            st.native = address <= 0;
            st.address = address < 0 ? -address : address;
            st.numberOfCalls = it->second.GetNumberOfCalls();
            st.executionTime = it->second.GetExecutionTime();
            stats.push_back(st);
        }
    }

    return stats;
}

<|endoftext|>


<commit_before>#include "ServerRepo.hpp"
#include "../inanity/data/sqlite.hpp"
#include "../inanity/MemoryFile.hpp"
#include "../inanity/PartFile.hpp"
#include "../inanity/StreamReader.hpp"
#include "../inanity/StreamWriter.hpp"
#include "../inanity/Exception.hpp"

BEGIN_INANITY_OIL

/*

Push:

-> client revision (shortly big)
-> client upper revision (shortly big)
->{
    key size (shortly)
    key data[key size]
    value size (shortly)
    value data[value size]
    base revision (shortly big)
  }
  0 (shortly)

<- total number of changed unique keys since client revision (shortly big)
<- pre-push revision (shortly big)
<- 1-based index of write[number of successful writes] (shortly)
<- 0 (shortly)
<-{
    key size (shortly)
    key data[key size]
    value size (shortly)
    value data[value size]
    revision (shortly big)
  }
  0 (shortly)
<- new client revision (shortly big)

*/

ServerRepo::ServerRepo(const char* fileName)
: Repo(fileName)
{
	BEGIN_TRY();

	// check format version
	CheckAppVersion(serverRepoAppVersion);

	// create table revs
	if(sqlite3_exec(*db,
		"CREATE TABLE IF NOT EXISTS revs ("
		"rev INTEGER PRIMARY KEY AUTOINCREMENT, "
		"key BLOB NOT NULL, "
		"value BLOB NOT NULL)",
		0, 0, 0) != SQLITE_OK)
		THROW_SECONDARY("Can't create table revs", db->Error());
	// create index revs_key_rev
	if(sqlite3_exec(*db,
		"CREATE INDEX IF NOT EXISTS revs_key_rev ON revs (key, rev)",
		0, 0, 0) != SQLITE_OK)
		THROW_SECONDARY("Can't create index revs_key_rev", db->Error());

	// create statements
	stmtGetMaxRevision = db->CreateStatement("SELECT MAX(rev) FROM revs");
	stmtCheckConflict = db->CreateStatement("SELECT COUNT(*) FROM revs WHERE key = ?1 AND rev > ?2");
	stmtWrite = db->CreateStatement("INSERT INTO revs (key, value) VALUES (?1, ?2)");
#define SUBQUERY "SELECT key, MAX(rev) AS maxrev FROM revs GROUP BY key"
	stmtPull = db->CreateStatement(
		"SELECT rev, key, value FROM revs NATURAL JOIN ("
		SUBQUERY
		") WHERE rev > ?1 AND rev <= ?2 AND rev = maxrev ORDER BY rev LIMIT ?3");
	stmtGetWeakRevision = db->CreateStatement(
		"SELECT rev FROM revs NATURAL JOIN ("
		SUBQUERY
		") WHERE rev > ?1 AND rev <= ?2 AND rev = maxrev ORDER BY rev LIMIT 1");
#undef SUBQUERY
	stmtPullTotalSize = db->CreateStatement(
		"SELECT COUNT(DISTINCT key) FROM revs WHERE rev > ?1");

	END_TRY("Can't create server repo");
}

long long ServerRepo::GetMaxRevision()
{
	Data::SqliteQuery query(stmtGetMaxRevision);

	if(stmtGetMaxRevision->Step() != SQLITE_ROW)
		THROW_SECONDARY("Can't get max revision", db->Error());

	return stmtGetMaxRevision->ColumnInt64(0);
}

void ServerRepo::WriteManifest(StreamWriter* writer)
{
	BEGIN_TRY();

	// write protocol magic
	writer->Write(protocolMagic, sizeof(protocolMagic));
	// write protocol version
	writer->WriteShortly(1);
	// write constraints
	writer->WriteShortly(maxKeySize);
	writer->WriteShortly(maxValueSize);
	writer->WriteShortly(maxPushKeysCount);
	writer->WriteShortly(maxPushTotalSize);
	writer->WriteShortly(maxPullKeysCount);
	writer->WriteShortly(maxPullTotalSize);

	END_TRY("Can't write server repo manifest");
}

void ServerRepo::Sync(StreamReader* reader, StreamWriter* writer)
{
	BEGIN_TRY();

	// push & pull should be in a one transaction
	Data::SqliteTransaction transaction(db);

	// get initial revision
	long long prePushRevision = GetMaxRevision();

	// read client revision
	long long clientRevision = reader->ReadShortlyBig();
	// read client upper revision
	long long clientUpperRevision = reader->ReadShortlyBig();
	// correct upper revision
	if(clientUpperRevision == 0 || clientUpperRevision > prePushRevision)
		clientUpperRevision = prePushRevision;

	// determine total size of the pull
	{
		Data::SqliteQuery queryPullTotalSize(stmtPullTotalSize);

		stmtPullTotalSize->Bind(1, clientRevision);
		if(stmtPullTotalSize->Step() != SQLITE_ROW)
			THROW("Can't determine total pull size");

		// output total size of pull
		long long pullTotalSize = stmtPullTotalSize->ColumnInt64(0);
		writer->WriteShortlyBig(pullTotalSize);
	}

	//*** push

	void* key = keyBufferFile->GetData();
	void* value = valueBufferFile->GetData();

	size_t totalPushSize = 0;

	// output pre-push revision
	writer->WriteShortlyBig(prePushRevision);

	// loop for push keys
	for(size_t i = 0; ; ++i)
	{
		// read key size
		size_t keySize = reader->ReadShortly();
		// if it's 0, it's signal of end
		if(!keySize)
			break;

		// check number of keys
		if(i >= (size_t)maxPushKeysCount)
			THROW("Too many keys");

		// check key size
		if(keySize > maxKeySize)
			THROW("Too big key");

		// read key value
		reader->Read(key, keySize);

		// read value
		size_t valueSize = reader->ReadShortly();
		if(valueSize > maxValueSize)
			THROW("Too big value");
		if(valueSize)
			reader->Read(value, valueSize);

		// read base revision
		long long baseRevision = reader->ReadShortlyBig();

		// check total push size
		totalPushSize += valueSize;
		if(totalPushSize > maxPushTotalSize)
			THROW("Too big push total size");

		// check for conflict
		Data::SqliteQuery queryCheckConflict(stmtCheckConflict);
		ptr<File> keyFile = NEW(PartFile(keyBufferFile, key, keySize));
		stmtCheckConflict->Bind(1, keyFile);
		stmtCheckConflict->Bind(2, baseRevision);
		if(stmtCheckConflict->Step() != SQLITE_ROW)
			THROW_SECONDARY("Can't check conflict", db->Error());
		if(stmtCheckConflict->ColumnInt(0) > 0)
			// conflict, skip this
			continue;

		// do write
		Data::SqliteQuery queryWrite(stmtWrite);
		stmtWrite->Bind(1, keyFile);
		stmtWrite->Bind(2, NEW(PartFile(valueBufferFile, value, valueSize)));
		if(stmtWrite->Step() != SQLITE_DONE)
			THROW_SECONDARY("Can't do write", db->Error());

		// output write
		writer->WriteShortly(i + 1);
	}

	// ensure request is over
	reader->ReadEnd();

	// output zero
	writer->WriteShortly(0);

	//*** pull
	Data::SqliteQuery queryPull(stmtPull);

	stmtPull->Bind(1, clientRevision);
	stmtPull->Bind(2, clientUpperRevision);
	stmtPull->Bind(3, maxPullKeysCount);

	long long lastKnownClientRevision = clientRevision;

	size_t totalPullSize = 0;
	bool done = false;
	while(!done)
	{
		switch(stmtPull->Step())
		{
		case SQLITE_ROW:
			// check total pull size
			totalPullSize += stmtPull->ColumnBlobSize(2);
			if(totalPullSize > maxPullTotalSize)
			{
				// stop pull, that's enough
				done = true;
				break;
			}

			// output key
			{
				ptr<File> key = stmtPull->ColumnBlob(1);
				size_t keySize = key->GetSize();
				writer->WriteShortly(keySize);
				writer->Write(key->GetData(), keySize);
			}
			// output value
			{
				ptr<File> value = stmtPull->ColumnBlob(2);
				size_t valueSize = value->GetSize();
				writer->WriteShortly(valueSize);
				writer->Write(value->GetData(), valueSize);
			}
			{
				// output revision
				long long revision = stmtPull->ColumnInt64(0);
				writer->WriteShortlyBig(revision);

				// remember last revision
				lastKnownClientRevision = revision;
			}

			break;

		case SQLITE_DONE:
			// end, no more keys
			done = true;
			break;

		default:
			// some error
			THROW_SECONDARY("Error pulling changes", db->Error());
		}
	}

	// write final zero
	writer->WriteShortly(0);

	// write new client revision
	{
		Data::SqliteQuery query(stmtGetWeakRevision);

		stmtGetWeakRevision->Bind(1, lastKnownClientRevision);
		stmtGetWeakRevision->Bind(2, clientUpperRevision);

		long long newClientRevision;
		switch(stmtGetWeakRevision->Step())
		{
		case SQLITE_ROW:
			newClientRevision = stmtGetWeakRevision->ColumnInt64(0) - 1;
			break;
		case SQLITE_DONE:
			newClientRevision = clientUpperRevision;
			break;
		default:
			THROW_SECONDARY("Error getting new client revision", db->Error());
		}

		writer->WriteShortlyBig(newClientRevision);
	}

	// commit transaction
	transaction.Commit();

	END_TRY("Can't sync server repo");
}

END_INANITY_OIL
<commit_msg>fix protocol version from "magic number"<commit_after>#include "ServerRepo.hpp"
#include "../inanity/data/sqlite.hpp"
#include "../inanity/MemoryFile.hpp"
#include "../inanity/PartFile.hpp"
#include "../inanity/StreamReader.hpp"
#include "../inanity/StreamWriter.hpp"
#include "../inanity/Exception.hpp"

BEGIN_INANITY_OIL

/*

Push:

-> client revision (shortly big)
-> client upper revision (shortly big)
->{
    key size (shortly)
    key data[key size]
    value size (shortly)
    value data[value size]
    base revision (shortly big)
  }
  0 (shortly)

<- total number of changed unique keys since client revision (shortly big)
<- pre-push revision (shortly big)
<- 1-based index of write[number of successful writes] (shortly)
<- 0 (shortly)
<-{
    key size (shortly)
    key data[key size]
    value size (shortly)
    value data[value size]
    revision (shortly big)
  }
  0 (shortly)
<- new client revision (shortly big)

*/

ServerRepo::ServerRepo(const char* fileName)
: Repo(fileName)
{
	BEGIN_TRY();

	// check format version
	CheckAppVersion(serverRepoAppVersion);

	// create table revs
	if(sqlite3_exec(*db,
		"CREATE TABLE IF NOT EXISTS revs ("
		"rev INTEGER PRIMARY KEY AUTOINCREMENT, "
		"key BLOB NOT NULL, "
		"value BLOB NOT NULL)",
		0, 0, 0) != SQLITE_OK)
		THROW_SECONDARY("Can't create table revs", db->Error());
	// create index revs_key_rev
	if(sqlite3_exec(*db,
		"CREATE INDEX IF NOT EXISTS revs_key_rev ON revs (key, rev)",
		0, 0, 0) != SQLITE_OK)
		THROW_SECONDARY("Can't create index revs_key_rev", db->Error());

	// create statements
	stmtGetMaxRevision = db->CreateStatement("SELECT MAX(rev) FROM revs");
	stmtCheckConflict = db->CreateStatement("SELECT COUNT(*) FROM revs WHERE key = ?1 AND rev > ?2");
	stmtWrite = db->CreateStatement("INSERT INTO revs (key, value) VALUES (?1, ?2)");
#define SUBQUERY "SELECT key, MAX(rev) AS maxrev FROM revs GROUP BY key"
	stmtPull = db->CreateStatement(
		"SELECT rev, key, value FROM revs NATURAL JOIN ("
		SUBQUERY
		") WHERE rev > ?1 AND rev <= ?2 AND rev = maxrev ORDER BY rev LIMIT ?3");
	stmtGetWeakRevision = db->CreateStatement(
		"SELECT rev FROM revs NATURAL JOIN ("
		SUBQUERY
		") WHERE rev > ?1 AND rev <= ?2 AND rev = maxrev ORDER BY rev LIMIT 1");
#undef SUBQUERY
	stmtPullTotalSize = db->CreateStatement(
		"SELECT COUNT(DISTINCT key) FROM revs WHERE rev > ?1");

	END_TRY("Can't create server repo");
}

long long ServerRepo::GetMaxRevision()
{
	Data::SqliteQuery query(stmtGetMaxRevision);

	if(stmtGetMaxRevision->Step() != SQLITE_ROW)
		THROW_SECONDARY("Can't get max revision", db->Error());

	return stmtGetMaxRevision->ColumnInt64(0);
}

void ServerRepo::WriteManifest(StreamWriter* writer)
{
	BEGIN_TRY();

	// write protocol magic
	writer->Write(protocolMagic, sizeof(protocolMagic));
	// write protocol version
	writer->WriteShortly(protocolVersion);
	// write constraints
	writer->WriteShortly(maxKeySize);
	writer->WriteShortly(maxValueSize);
	writer->WriteShortly(maxPushKeysCount);
	writer->WriteShortly(maxPushTotalSize);
	writer->WriteShortly(maxPullKeysCount);
	writer->WriteShortly(maxPullTotalSize);

	END_TRY("Can't write server repo manifest");
}

void ServerRepo::Sync(StreamReader* reader, StreamWriter* writer)
{
	BEGIN_TRY();

	// push & pull should be in a one transaction
	Data::SqliteTransaction transaction(db);

	// get initial revision
	long long prePushRevision = GetMaxRevision();

	// read client revision
	long long clientRevision = reader->ReadShortlyBig();
	// read client upper revision
	long long clientUpperRevision = reader->ReadShortlyBig();
	// correct upper revision
	if(clientUpperRevision == 0 || clientUpperRevision > prePushRevision)
		clientUpperRevision = prePushRevision;

	// determine total size of the pull
	{
		Data::SqliteQuery queryPullTotalSize(stmtPullTotalSize);

		stmtPullTotalSize->Bind(1, clientRevision);
		if(stmtPullTotalSize->Step() != SQLITE_ROW)
			THROW("Can't determine total pull size");

		// output total size of pull
		long long pullTotalSize = stmtPullTotalSize->ColumnInt64(0);
		writer->WriteShortlyBig(pullTotalSize);
	}

	//*** push

	void* key = keyBufferFile->GetData();
	void* value = valueBufferFile->GetData();

	size_t totalPushSize = 0;

	// output pre-push revision
	writer->WriteShortlyBig(prePushRevision);

	// loop for push keys
	for(size_t i = 0; ; ++i)
	{
		// read key size
		size_t keySize = reader->ReadShortly();
		// if it's 0, it's signal of end
		if(!keySize)
			break;

		// check number of keys
		if(i >= (size_t)maxPushKeysCount)
			THROW("Too many keys");

		// check key size
		if(keySize > maxKeySize)
			THROW("Too big key");

		// read key value
		reader->Read(key, keySize);

		// read value
		size_t valueSize = reader->ReadShortly();
		if(valueSize > maxValueSize)
			THROW("Too big value");
		if(valueSize)
			reader->Read(value, valueSize);

		// read base revision
		long long baseRevision = reader->ReadShortlyBig();

		// check total push size
		totalPushSize += valueSize;
		if(totalPushSize > maxPushTotalSize)
			THROW("Too big push total size");

		// check for conflict
		Data::SqliteQuery queryCheckConflict(stmtCheckConflict);
		ptr<File> keyFile = NEW(PartFile(keyBufferFile, key, keySize));
		stmtCheckConflict->Bind(1, keyFile);
		stmtCheckConflict->Bind(2, baseRevision);
		if(stmtCheckConflict->Step() != SQLITE_ROW)
			THROW_SECONDARY("Can't check conflict", db->Error());
		if(stmtCheckConflict->ColumnInt(0) > 0)
			// conflict, skip this
			continue;

		// do write
		Data::SqliteQuery queryWrite(stmtWrite);
		stmtWrite->Bind(1, keyFile);
		stmtWrite->Bind(2, NEW(PartFile(valueBufferFile, value, valueSize)));
		if(stmtWrite->Step() != SQLITE_DONE)
			THROW_SECONDARY("Can't do write", db->Error());

		// output write
		writer->WriteShortly(i + 1);
	}

	// ensure request is over
	reader->ReadEnd();

	// output zero
	writer->WriteShortly(0);

	//*** pull
	Data::SqliteQuery queryPull(stmtPull);

	stmtPull->Bind(1, clientRevision);
	stmtPull->Bind(2, clientUpperRevision);
	stmtPull->Bind(3, maxPullKeysCount);

	long long lastKnownClientRevision = clientRevision;

	size_t totalPullSize = 0;
	bool done = false;
	while(!done)
	{
		switch(stmtPull->Step())
		{
		case SQLITE_ROW:
			// check total pull size
			totalPullSize += stmtPull->ColumnBlobSize(2);
			if(totalPullSize > maxPullTotalSize)
			{
				// stop pull, that's enough
				done = true;
				break;
			}

			// output key
			{
				ptr<File> key = stmtPull->ColumnBlob(1);
				size_t keySize = key->GetSize();
				writer->WriteShortly(keySize);
				writer->Write(key->GetData(), keySize);
			}
			// output value
			{
				ptr<File> value = stmtPull->ColumnBlob(2);
				size_t valueSize = value->GetSize();
				writer->WriteShortly(valueSize);
				writer->Write(value->GetData(), valueSize);
			}
			{
				// output revision
				long long revision = stmtPull->ColumnInt64(0);
				writer->WriteShortlyBig(revision);

				// remember last revision
				lastKnownClientRevision = revision;
			}

			break;

		case SQLITE_DONE:
			// end, no more keys
			done = true;
			break;

		default:
			// some error
			THROW_SECONDARY("Error pulling changes", db->Error());
		}
	}

	// write final zero
	writer->WriteShortly(0);

	// write new client revision
	{
		Data::SqliteQuery query(stmtGetWeakRevision);

		stmtGetWeakRevision->Bind(1, lastKnownClientRevision);
		stmtGetWeakRevision->Bind(2, clientUpperRevision);

		long long newClientRevision;
		switch(stmtGetWeakRevision->Step())
		{
		case SQLITE_ROW:
			newClientRevision = stmtGetWeakRevision->ColumnInt64(0) - 1;
			break;
		case SQLITE_DONE:
			newClientRevision = clientUpperRevision;
			break;
		default:
			THROW_SECONDARY("Error getting new client revision", db->Error());
		}

		writer->WriteShortlyBig(newClientRevision);
	}

	// commit transaction
	transaction.Commit();

	END_TRY("Can't sync server repo");
}

END_INANITY_OIL
<|endoftext|>