<commit_before>/*BEGIN_LEGAL
Intel Open Source License

Copyright (c) 2002-2013 Intel Corporation. All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are
met:

Redistributions of source code must retain the above copyright notice,
this list of conditions and the following disclaimer.  Redistributions
in binary form must reproduce the above copyright notice, this list of
conditions and the following disclaimer in the documentation and/or
other materials provided with the distribution.  Neither the name of
the Intel Corporation nor the names of its contributors may be used to
endorse or promote products derived from this software without
specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE INTEL OR
ITS CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
END_LEGAL */

#include "pin.H"
#include <iostream>
#include <cstdio>
#include <cstring>
#include <set>
#include "memlist.h"
#include "memoryalloc.h"
#include "stats.h"
using namespace std;

// Prototypes ; TODO: Move to Seperate Header file?
void HookFree(IMG img);
void HookMalloc(IMG img);
void HookCalloc(IMG img);
void HookRealloc(IMG img);

typedef VOID* (*FP_MALLOC)(size_t);
typedef void (*FP_FREE)(void*);
typedef void* (*FP_CALLOC)(size_t, size_t);
typedef void* (*FP_REALLOC)(void*, size_t);

bool hasEnding (std::string const &fullString, std::string const &ending);

bool inMain = false;
FILE * trace;
MemList ml;
Stats stats;

// Print a memory read record
VOID RecordHeapMemRead(VOID * ip, VOID * addr) {
    int rtn = ml.containsAddress(addr);
    if(rtn == ERR_IN_FENCE) {
        fprintf(trace,"##########BAD WRITE: %p \n", addr);
        cout << "BAD READ" << endl;
        stats.incInvalidReadCount();
    }
    //printf("heap read: %p\n", addr);
}

// Print a memory write record
VOID RecordHeapMemWrite(VOID * ip, VOID * addr) {
    int rtn = ml.containsAddress(addr);
    if(rtn == ERR_IN_FENCE) {
        fprintf(trace,"##########BAD WRITE: %p \n", addr);
        cout << "BAD WRITE" << endl;
        stats.incInvalidWriteCount();
    }
    //printf("heap write: %p\n", addr);
}

bool hasEnding (std::string const &fullString, std::string const &ending) {
    if (fullString.length() >= ending.length()) {
        return (0 == fullString.compare (fullString.length() - ending.length(), ending.length(), ending));
    } else {
        return false;
    }
}

VOID RecordStackMemRead(VOID * ip, VOID * addr) {
    //printf("stack read: %p\n", addr);
}

// Print a memory write record
VOID RecordStackMemWrite(VOID * ip, VOID * addr) {
    //printf("stack write: %p\n", addr);
}

// Set global inMain = true
VOID SetInMain(void){
	inMain = true;
	cout << "Main execution started" << endl;
}

// Is called for every instruction and instruments reads and writes
VOID Instruction(INS ins, VOID *v) {
    if(!inMain)
        return;
    // Instruments memory accesses using a predicated call, i.e.
    // the instrumentation is called iff the instruction will actually be executed.
    //
    // On the IA-32 and Intel(R) 64 architectures conditional moves and REP
    // prefixed instructions appear as predicated instructions in Pin.
    UINT32 memOperands = INS_MemoryOperandCount(ins);

    // Iterate over each memory operand of the instruction.
    for (UINT32 memOp = 0; memOp < memOperands; memOp++) {
        //if (INS_MemoryOperandIsRead(ins, memOp)) {
        if (INS_MemoryOperandIsRead(ins, memOp) && !INS_IsStackRead(ins)) {
            INS_InsertPredicatedCall(
                ins, IPOINT_BEFORE, (AFUNPTR)RecordHeapMemRead,
                IARG_INST_PTR,
                IARG_MEMORYOP_EA, memOp,
                IARG_END);
        } else if (INS_MemoryOperandIsRead(ins, memOp)) {
            INS_InsertPredicatedCall(
                ins, IPOINT_BEFORE, (AFUNPTR)RecordStackMemRead,
                IARG_INST_PTR,
                IARG_MEMORYOP_EA, memOp,
                IARG_END);
        }
        // Note that in some architectures a single memory operand can be
        // both read and written (for instance incl (%eax) on IA-32)
        // In that case we instrument it once for read and once for write.
        //if (INS_MemoryOperandIsWritten(ins, memOp)) {
        if (INS_MemoryOperandIsWritten(ins, memOp) && !INS_IsStackWrite(ins)) {
            INS_InsertPredicatedCall(
                ins, IPOINT_BEFORE, (AFUNPTR)RecordHeapMemWrite,
                IARG_INST_PTR,
                IARG_MEMORYOP_EA, memOp,
                IARG_END);
        } else if (INS_MemoryOperandIsWritten(ins, memOp)) {
            INS_InsertPredicatedCall(
                ins, IPOINT_BEFORE, (AFUNPTR)RecordStackMemWrite,
                IARG_INST_PTR,
                IARG_MEMORYOP_EA, memOp,
                IARG_END);
        }
    }
}

VOID Fini(INT32 code, VOID *v) {
    fprintf(trace, "#eof \n");
    // Display the stats
    stats.displayResults(ml, trace);
    fclose(trace);
}

// This is the replacement routine.
VOID* NewMalloc(FP_MALLOC orgFuncptr, UINT32 arg0, ADDRINT returnIp) {
    // Call the relocated entry point of the original (replaced) routine.
    void* v = orgFuncptr(arg0 + (2 * DEFAULT_FENCE_SIZE));
    stats.incMallocCount();
    MemoryAlloc ma = ml.add(v, arg0, DEFAULT_FENCE_SIZE);
    fprintf(trace, "ADDED: %p %d \n", v, DEFAULT_FENCE_SIZE);
    return ma.getAddress();
}

void* NewCalloc(FP_CALLOC libc_calloc, UINT32 arg0, UINT32 arg1, ADDRINT returnIp) {
    // Calculate the size in bytes
    size_t bytes = (arg0 * arg1);
    size_t totalSize = bytes + (2 + DEFAULT_FENCE_SIZE);
    // Figure out the correct amount of chunks
    size_t nmemb = totalSize / arg1;
    size_t nmembe = totalSize % arg1;
    nmembe = (nmembe > 0) ? 1 : 0;
    // Sum up the chunks
    nmemb += nmembe;
    // Allocate space using malloc
    void *ptr = libc_calloc(nmemb, arg1);
    // Create a new MemoryAlloc
    MemoryAlloc ma = ml.add(ptr, bytes, DEFAULT_FENCE_SIZE);
    // Add the information to the trace file
    fprintf(trace, "ADDED: %p %d \n", ptr, DEFAULT_FENCE_SIZE);
    stats.incMallocCount();
    return ma.getAddress();
}

void* NewRealloc(FP_REALLOC orgFuncptr, void* arg0, UINT32 arg1, ADDRINT returnIp) {
    printf("New Realloc was called\n");
    return arg0;
}

void NewFree(FP_FREE orgFuncptr, void* ptr, ADDRINT returnIp) {
    if(ptr != NULL) {
        // Check the MemList
        int index  = ml.containsAddress(ptr);
        if(index >= 0) {
            MemoryAlloc alloc = ml.get(index);
            // Remove the space
            orgFuncptr(alloc.getUnderflowFence());
            fprintf(trace, "Freed the memory @ address %p successfully.\n", alloc.getUnderflowFence());
            // Remove the element from the list
            ml.removeMatching(alloc);
            stats.incFreeCount();
        } else if(index == ERR_NOT_FOUND) {
            // This currently gets hit since a blacklist is being used
            fprintf(trace, "Address = %p not found. Bad address or stack address used.\n", ptr);
            stats.incInvalidFreeCount();
            /*
            // Dump all memory currently stored
            char buffer[1024];
            for(int i = 0; i < ml.size(); i++) {
                fprintf(trace, "%s\n", ml.get(i).toString(buffer, 1024));
            }
            */
        } else if(index == ERR_MID_CHUNK) {
            fprintf(trace, "Mid-chunk memory deallocation @ %p\n", ptr);
            stats.incMidFreeChunkCount();
        } else if(index == ERR_IN_FENCE) {
            fprintf(trace, "Memory Fence Hit @ %p\n", ptr);
            stats.incFenceHitCount();
        } else {
            // This should Never happen...
            fprintf(trace, "index = %d : Unable to deallocate the memory @ %p\n", index, ptr);
        }
    } else {
        fprintf(trace, "Attempted to free NULL\n");
        stats.incFreeNullCount();
    }
}

// Pin calls this function every time a new img is loaded.
// It is best to do probe replacement when the image is loaded,
// because only one thread knows about the image at this time.
VOID ImageLoad(IMG img, VOID *v) {

    if(IMG_IsMainExecutable(img)) {
		RTN rtn = RTN_FindByName(img, "main");
        RTN_Open(rtn);//Must open RTN API before examining instructions
		INS ins = RTN_InsHead(rtn);
        INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)SetInMain, IARG_END);
        RTN_Close(rtn);
    }

    if(!hasEnding(IMG_Name(img), "libc.so.6")) {
        return;
    }

    if(!inMain) {
        inMain = IMG_IsMainExecutable(img);
    }

    // Hook Functions
    HookMalloc(img);
    HookFree(img);
    HookCalloc(img);
    HookRealloc(img);
}

void HookFree(IMG img) {
    RTN rtn = RTN_FindByName(img, "free");
    if(RTN_Valid(rtn)) {
        cout << "Replacing free in " << IMG_Name(img) << endl;
        // Return type, cstype, function name, arguments...
        PROTO proto = PROTO_Allocate(PIN_PARG(void), CALLINGSTD_DEFAULT, "free", PIN_PARG(void*), PIN_PARG_END());
        // Replace free
        RTN_ReplaceSignature(rtn, AFUNPTR(NewFree),
            IARG_PROTOTYPE,
            proto,
            IARG_ORIG_FUNCPTR,
            IARG_FUNCARG_ENTRYPOINT_VALUE,
            0,
            IARG_RETURN_IP,
            IARG_END);
        // Free the function prototype
        PROTO_Free(proto);
    }
}

void HookMalloc(IMG img) {
    // See if malloc() is present in the image.  If so, replace it.
    RTN rtnMalloc = RTN_FindByName(img, "malloc");

    if(RTN_Valid(rtnMalloc)) {
        cout << "Replacing malloc in " << IMG_Name(img) << endl;
        PROTO proto_malloc = PROTO_Allocate(PIN_PARG(void *), CALLINGSTD_DEFAULT,
            "malloc", PIN_PARG(int), PIN_PARG_END());
        RTN_ReplaceSignature(rtnMalloc, AFUNPTR(NewMalloc),
            IARG_PROTOTYPE, proto_malloc,
            IARG_ORIG_FUNCPTR,
            IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
            IARG_RETURN_IP,
            IARG_END);
        // Free the function prototype.
        PROTO_Free(proto_malloc);
    }
}

void HookCalloc(IMG img) {
    // See if calloc() is present in the image.  If so, replace it.
    RTN rtn = RTN_FindByName(img, "calloc");
    if(RTN_Valid(rtn)) {
        cout << "Replacing calloc in " << IMG_Name(img) << endl;
        PROTO proto = PROTO_Allocate(PIN_PARG(void *), CALLINGSTD_DEFAULT,
            "calloc", PIN_PARG(size_t), PIN_PARG(size_t), PIN_PARG_END());
        RTN_ReplaceSignature(rtn, AFUNPTR(NewCalloc),
            IARG_PROTOTYPE, proto,
            IARG_ORIG_FUNCPTR,
            IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
            IARG_RETURN_IP,
            IARG_END);
        // Free the function prototype.
        PROTO_Free(proto);
    }
}

void HookRealloc(IMG img) {
    // See if calloc() is present in the image.  If so, replace it.
    RTN rtn = RTN_FindByName(img, "realloc");
    if(RTN_Valid(rtn)) {
        cout << "Replacing realloc in " << IMG_Name(img) << endl;
        PROTO proto = PROTO_Allocate(PIN_PARG(void *), CALLINGSTD_DEFAULT,
            "realloc", PIN_PARG(void*), PIN_PARG(size_t), PIN_PARG_END());
        RTN_ReplaceSignature(rtn, AFUNPTR(NewRealloc),
            IARG_PROTOTYPE, proto,
            IARG_ORIG_FUNCPTR,
            IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
            IARG_RETURN_IP,
            IARG_END);
        // Free the function prototype.
        PROTO_Free(proto);
    }
}

/* ===================================================================== */
/* Print Help Message                                                    */
/* ===================================================================== */

INT32 Usage() {
    cerr << "This tool demonstrates how to replace an original" << endl;
    cerr << " function with a custom function defined in the tool " << endl;
    cerr << " using probes.  The replacement function has a different " << endl;
    cerr << " signature from that of the original replaced function." << endl;
    cerr << endl << KNOB_BASE::StringKnobSummary() << endl;
    return -1;
}

/* ===================================================================== */
/* Main: Initialize and start Pin in Probe mode.                         */
/* ===================================================================== */

int main(INT32 argc, CHAR *argv[]) {
    // Initialize symbol processing
    PIN_InitSymbols();
    // Initialize pin
    if (PIN_Init(argc, argv))
        return Usage();
    // Open up the trace file
    trace = fopen("zzz.out", "w");
    // Register ImageLoad to be called when an image is loaded
    IMG_AddInstrumentFunction(ImageLoad, 0);
    INS_AddInstrumentFunction(Instruction, 0);
    PIN_AddFiniFunction(Fini, 0);
    PIN_StartProgram();
    return 0;
}
<commit_msg>Removed print statements from stdout.<commit_after>/*BEGIN_LEGAL
Intel Open Source License

Copyright (c) 2002-2013 Intel Corporation. All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are
met:

Redistributions of source code must retain the above copyright notice,
this list of conditions and the following disclaimer.  Redistributions
in binary form must reproduce the above copyright notice, this list of
conditions and the following disclaimer in the documentation and/or
other materials provided with the distribution.  Neither the name of
the Intel Corporation nor the names of its contributors may be used to
endorse or promote products derived from this software without
specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE INTEL OR
ITS CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
END_LEGAL */

#include "pin.H"
#include <iostream>
#include <cstdio>
#include <cstring>
#include <set>
#include "memlist.h"
#include "memoryalloc.h"
#include "stats.h"
using namespace std;

// Prototypes ; TODO: Move to Seperate Header file?
void HookFree(IMG img);
void HookMalloc(IMG img);
void HookCalloc(IMG img);
void HookRealloc(IMG img);

typedef VOID* (*FP_MALLOC)(size_t);
typedef void (*FP_FREE)(void*);
typedef void* (*FP_CALLOC)(size_t, size_t);
typedef void* (*FP_REALLOC)(void*, size_t);

bool hasEnding (std::string const &fullString, std::string const &ending);

bool inMain = false;
FILE * trace;
MemList ml;
Stats stats;

// Print a memory read record
VOID RecordHeapMemRead(VOID * ip, VOID * addr) {
    int rtn = ml.containsAddress(addr);
    if(rtn == ERR_IN_FENCE) {
        fprintf(trace,"##########BAD WRITE: %p \n", addr);
        stats.incInvalidReadCount();
    }
    //printf("heap read: %p\n", addr);
}

// Print a memory write record
VOID RecordHeapMemWrite(VOID * ip, VOID * addr) {
    int rtn = ml.containsAddress(addr);
    if(rtn == ERR_IN_FENCE) {
        fprintf(trace,"##########BAD WRITE: %p \n", addr);
        stats.incInvalidWriteCount();
    }
    //printf("heap write: %p\n", addr);
}

bool hasEnding (std::string const &fullString, std::string const &ending) {
    if (fullString.length() >= ending.length()) {
        return (0 == fullString.compare (fullString.length() - ending.length(), ending.length(), ending));
    } else {
        return false;
    }
}

VOID RecordStackMemRead(VOID * ip, VOID * addr) {
    //printf("stack read: %p\n", addr);
}

// Print a memory write record
VOID RecordStackMemWrite(VOID * ip, VOID * addr) {
    //printf("stack write: %p\n", addr);
}

// Set global inMain = true
VOID SetInMain(void){
	inMain = true;
    fprintf(trace, "Main execution started\n");
}

// Is called for every instruction and instruments reads and writes
VOID Instruction(INS ins, VOID *v) {
    if(!inMain)
        return;
    // Instruments memory accesses using a predicated call, i.e.
    // the instrumentation is called iff the instruction will actually be executed.
    //
    // On the IA-32 and Intel(R) 64 architectures conditional moves and REP
    // prefixed instructions appear as predicated instructions in Pin.
    UINT32 memOperands = INS_MemoryOperandCount(ins);

    // Iterate over each memory operand of the instruction.
    for (UINT32 memOp = 0; memOp < memOperands; memOp++) {
        //if (INS_MemoryOperandIsRead(ins, memOp)) {
        if (INS_MemoryOperandIsRead(ins, memOp) && !INS_IsStackRead(ins)) {
            INS_InsertPredicatedCall(
                ins, IPOINT_BEFORE, (AFUNPTR)RecordHeapMemRead,
                IARG_INST_PTR,
                IARG_MEMORYOP_EA, memOp,
                IARG_END);
        } else if (INS_MemoryOperandIsRead(ins, memOp)) {
            INS_InsertPredicatedCall(
                ins, IPOINT_BEFORE, (AFUNPTR)RecordStackMemRead,
                IARG_INST_PTR,
                IARG_MEMORYOP_EA, memOp,
                IARG_END);
        }
        // Note that in some architectures a single memory operand can be
        // both read and written (for instance incl (%eax) on IA-32)
        // In that case we instrument it once for read and once for write.
        //if (INS_MemoryOperandIsWritten(ins, memOp)) {
        if (INS_MemoryOperandIsWritten(ins, memOp) && !INS_IsStackWrite(ins)) {
            INS_InsertPredicatedCall(
                ins, IPOINT_BEFORE, (AFUNPTR)RecordHeapMemWrite,
                IARG_INST_PTR,
                IARG_MEMORYOP_EA, memOp,
                IARG_END);
        } else if (INS_MemoryOperandIsWritten(ins, memOp)) {
            INS_InsertPredicatedCall(
                ins, IPOINT_BEFORE, (AFUNPTR)RecordStackMemWrite,
                IARG_INST_PTR,
                IARG_MEMORYOP_EA, memOp,
                IARG_END);
        }
    }
}

VOID Fini(INT32 code, VOID *v) {
    fprintf(trace, "#eof \n");
    // Display the stats
    stats.displayResults(ml, trace);
    fclose(trace);
}

// This is the replacement routine.
VOID* NewMalloc(FP_MALLOC orgFuncptr, UINT32 arg0, ADDRINT returnIp) {
    // Call the relocated entry point of the original (replaced) routine.
    void* v = orgFuncptr(arg0 + (2 * DEFAULT_FENCE_SIZE));
    stats.incMallocCount();
    MemoryAlloc ma = ml.add(v, arg0, DEFAULT_FENCE_SIZE);
    fprintf(trace, "ADDED: %p %d \n", v, DEFAULT_FENCE_SIZE);
    return ma.getAddress();
}

void* NewCalloc(FP_CALLOC libc_calloc, UINT32 arg0, UINT32 arg1, ADDRINT returnIp) {
    // Calculate the size in bytes
    size_t bytes = (arg0 * arg1);
    size_t totalSize = bytes + (2 + DEFAULT_FENCE_SIZE);
    // Figure out the correct amount of chunks
    size_t nmemb = totalSize / arg1;
    size_t nmembe = totalSize % arg1;
    nmembe = (nmembe > 0) ? 1 : 0;
    // Sum up the chunks
    nmemb += nmembe;
    // Allocate space using malloc
    void *ptr = libc_calloc(nmemb, arg1);
    // Create a new MemoryAlloc
    MemoryAlloc ma = ml.add(ptr, bytes, DEFAULT_FENCE_SIZE);
    // Add the information to the trace file
    fprintf(trace, "ADDED: %p %d \n", ptr, DEFAULT_FENCE_SIZE);
    stats.incMallocCount();
    return ma.getAddress();
}

void* NewRealloc(FP_REALLOC orgFuncptr, void* arg0, UINT32 arg1, ADDRINT returnIp) {
    printf("New Realloc was called\n");
    return arg0;
}

void NewFree(FP_FREE orgFuncptr, void* ptr, ADDRINT returnIp) {
    if(ptr != NULL) {
        // Check the MemList
        int index  = ml.containsAddress(ptr);
        if(index >= 0) {
            MemoryAlloc alloc = ml.get(index);
            // Remove the space
            orgFuncptr(alloc.getUnderflowFence());
            fprintf(trace, "Freed the memory @ address %p successfully.\n", alloc.getUnderflowFence());
            // Remove the element from the list
            ml.removeMatching(alloc);
            stats.incFreeCount();
        } else if(index == ERR_NOT_FOUND) {
            // This currently gets hit since a blacklist is being used
            fprintf(trace, "Address = %p not found. Bad address or stack address used.\n", ptr);
            stats.incInvalidFreeCount();
            /*
            // Dump all memory currently stored
            char buffer[1024];
            for(int i = 0; i < ml.size(); i++) {
                fprintf(trace, "%s\n", ml.get(i).toString(buffer, 1024));
            }
            */
        } else if(index == ERR_MID_CHUNK) {
            fprintf(trace, "Mid-chunk memory deallocation @ %p\n", ptr);
            stats.incMidFreeChunkCount();
        } else if(index == ERR_IN_FENCE) {
            fprintf(trace, "Memory Fence Hit @ %p\n", ptr);
            stats.incFenceHitCount();
        } else {
            // This should Never happen...
            fprintf(trace, "index = %d : Unable to deallocate the memory @ %p\n", index, ptr);
        }
    } else {
        fprintf(trace, "Attempted to free NULL\n");
        stats.incFreeNullCount();
    }
}

// Pin calls this function every time a new img is loaded.
// It is best to do probe replacement when the image is loaded,
// because only one thread knows about the image at this time.
VOID ImageLoad(IMG img, VOID *v) {

    if(IMG_IsMainExecutable(img)) {
		RTN rtn = RTN_FindByName(img, "main");
        RTN_Open(rtn);//Must open RTN API before examining instructions
		INS ins = RTN_InsHead(rtn);
        INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)SetInMain, IARG_END);
        RTN_Close(rtn);
    }

    if(!hasEnding(IMG_Name(img), "libc.so.6")) {
        return;
    }

    if(!inMain) {
        inMain = IMG_IsMainExecutable(img);
    }

    // Hook Functions
    HookMalloc(img);
    HookFree(img);
    HookCalloc(img);
    HookRealloc(img);
}

void HookFree(IMG img) {
    RTN rtn = RTN_FindByName(img, "free");
    if(RTN_Valid(rtn)) {
        fprintf(trace, "Replacing free in %s\n", IMG_Name(img).c_str());
        // Return type, cstype, function name, arguments...
        PROTO proto = PROTO_Allocate(PIN_PARG(void), CALLINGSTD_DEFAULT, "free", PIN_PARG(void*), PIN_PARG_END());
        // Replace free
        RTN_ReplaceSignature(rtn, AFUNPTR(NewFree),
            IARG_PROTOTYPE,
            proto,
            IARG_ORIG_FUNCPTR,
            IARG_FUNCARG_ENTRYPOINT_VALUE,
            0,
            IARG_RETURN_IP,
            IARG_END);
        // Free the function prototype
        PROTO_Free(proto);
    }
}

void HookMalloc(IMG img) {
    // See if malloc() is present in the image.  If so, replace it.
    RTN rtnMalloc = RTN_FindByName(img, "malloc");

    if(RTN_Valid(rtnMalloc)) {
        fprintf(trace, "Replacing malloc in %s\n", IMG_Name(img).c_str());
        PROTO proto_malloc = PROTO_Allocate(PIN_PARG(void *), CALLINGSTD_DEFAULT,
            "malloc", PIN_PARG(int), PIN_PARG_END());
        RTN_ReplaceSignature(rtnMalloc, AFUNPTR(NewMalloc),
            IARG_PROTOTYPE, proto_malloc,
            IARG_ORIG_FUNCPTR,
            IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
            IARG_RETURN_IP,
            IARG_END);
        // Free the function prototype.
        PROTO_Free(proto_malloc);
    }
}

void HookCalloc(IMG img) {
    // See if calloc() is present in the image.  If so, replace it.
    RTN rtn = RTN_FindByName(img, "calloc");
    if(RTN_Valid(rtn)) {
        fprintf(trace, "Replacing calloc in %s\n", IMG_Name(img).c_str());
        PROTO proto = PROTO_Allocate(PIN_PARG(void *), CALLINGSTD_DEFAULT,
            "calloc", PIN_PARG(size_t), PIN_PARG(size_t), PIN_PARG_END());
        RTN_ReplaceSignature(rtn, AFUNPTR(NewCalloc),
            IARG_PROTOTYPE, proto,
            IARG_ORIG_FUNCPTR,
            IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
            IARG_RETURN_IP,
            IARG_END);
        // Free the function prototype.
        PROTO_Free(proto);
    }
}

void HookRealloc(IMG img) {
    // See if calloc() is present in the image.  If so, replace it.
    RTN rtn = RTN_FindByName(img, "realloc");
    if(RTN_Valid(rtn)) {
        fprintf(trace, "Replacing realloc in %s\n", IMG_Name(img).c_str());
        PROTO proto = PROTO_Allocate(PIN_PARG(void *), CALLINGSTD_DEFAULT,
            "realloc", PIN_PARG(void*), PIN_PARG(size_t), PIN_PARG_END());
        RTN_ReplaceSignature(rtn, AFUNPTR(NewRealloc),
            IARG_PROTOTYPE, proto,
            IARG_ORIG_FUNCPTR,
            IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
            IARG_RETURN_IP,
            IARG_END);
        // Free the function prototype.
        PROTO_Free(proto);
    }
}

/* ===================================================================== */
/* Print Help Message                                                    */
/* ===================================================================== */

INT32 Usage() {
    cerr << "This tool demonstrates how to replace an original" << endl;
    cerr << " function with a custom function defined in the tool " << endl;
    cerr << " using probes.  The replacement function has a different " << endl;
    cerr << " signature from that of the original replaced function." << endl;
    cerr << endl << KNOB_BASE::StringKnobSummary() << endl;
    return -1;
}

/* ===================================================================== */
/* Main: Initialize and start Pin in Probe mode.                         */
/* ===================================================================== */

int main(INT32 argc, CHAR *argv[]) {
    // Initialize symbol processing
    PIN_InitSymbols();
    // Initialize pin
    if (PIN_Init(argc, argv))
        return Usage();
    // Open up the trace file
    trace = fopen("zzz.out", "w");
    // Register ImageLoad to be called when an image is loaded
    IMG_AddInstrumentFunction(ImageLoad, 0);
    INS_AddInstrumentFunction(Instruction, 0);
    PIN_AddFiniFunction(Fini, 0);
    PIN_StartProgram();
    return 0;
}
<|endoftext|>


<commit_before>#include "stats.h"

Stats::Stats() {
	this->reset();
}

void Stats::reset() {
	/* Allocation Stats */
	this->allocCount = 0;
	this->mallocCount = 0;
	this->callocCount = 0;
	this->reallocCount = 0;
	/* Deallocation Stats */
	this->freeCount = 0;
	/* Invalid Free Stats*/
	this->invalidFreeCount = 0;
	this->midFreeChunkCount = 0;
	this->freeNullCount = 0;
	/* Invalid Read Stats */
	this->invalidReadCount = 0;
	this->readFenceOverflow = 0;
	this->readFenceUnderflow = 0;
	/* Invalid Write Stats */
	this->invalidWriteCount = 0;
	this->writeFenceOverflow = 0;
	this->writeFenceUnderflow = 0;
	/* Stack Smashing Stats */
	this->invalidReturnCount = 0;
}
/* Getters */
unsigned int Stats::getAllocCount() {
	return this->allocCount;
}

unsigned int Stats::getMallocCount() {
	return this->mallocCount;
}

unsigned int Stats::getCallocCount() {
	return this->callocCount;
}

unsigned int Stats::getReallocCount() {
	return this->reallocCount;
}

unsigned int Stats::getFreeCount() {
	return this->freeCount;
}
unsigned int Stats::getInvalidReadCount() {
	return this->invalidReadCount;
}

unsigned int Stats::getReadFenceOverflow() {
	return this->readFenceOverflow;
}

unsigned int Stats::getReadFenceUnderflow() {
	return this->readFenceUnderflow;
}

unsigned int Stats::getInvalidWriteCount() {
	return this->invalidWriteCount;
}

unsigned int Stats::getWriteFenceOverflow() {
	return this->writeFenceOverflow;
}

unsigned int Stats::getWriteFenceUnderflow() {
	return this->writeFenceUnderflow;
}

unsigned int Stats::getInvalidFreeCount() {
	return this->invalidFreeCount;
}

unsigned int Stats::getMidFreeChunkCount() {
	return this->midFreeChunkCount;
}

unsigned int Stats::getFreeNullCount() {
	return this->freeNullCount;
}

unsigned int Stats::getFreeNotFoundCount() {
	return this->freeNotFoundCount;
}

unsigned int Stats::getInvalidReturnCount() {
	return this->invalidReturnCount;
}

/* Setters */

void Stats::setAllocCount(unsigned int count) {
	this->allocCount = count;
}

void Stats::setMallocCount(unsigned int count) {
	this->allocCount = count;
}

void Stats::setCallocCount(unsigned int count) {
	this->allocCount = count;
}

void Stats::setReallocCount(unsigned int count) {
	this->allocCount = count;
}

void Stats::setFreeCount(unsigned int count) {
	this->freeCount = count;
}

void Stats::setInvalidReadCount(unsigned int count) {
	this->invalidReadCount = count;
}

void Stats::setReadFenceOverflow(unsigned int count) {
	this->readFenceOverflow = count;
}

void Stats::setReadFenceUnderflow(unsigned int count) {
	this->readFenceUnderflow = count;
}

void Stats::setInvalidWriteCount(unsigned int count) {
	this->invalidWriteCount = count;
}

void Stats::setWriteFenceOverflow(unsigned int count) {
	this->writeFenceOverflow = count;
}

void Stats::setWriteFenceUnderflow(unsigned int count) {
	this->writeFenceUnderflow = count;
}

void Stats::setInvalidFreeCount(unsigned int count) {
	this->invalidFreeCount = count;
}

void Stats::setMidFreeChunkCount(unsigned int count) {
	this->midFreeChunkCount = count;
}

void Stats::setFreeNullCount(unsigned int count) {
	this->freeNullCount = count;
}

void Stats::setInvalidReturnCount(unsigned int count) {
	this->invalidReturnCount = count;
}

void Stats::setFreeNotFoundCount(unsigned int count) {
	this->freeNotFoundCount = count;
}

/* Increment Count Functions */

void Stats::incAllocCount() {
	this->allocCount++;
}

void Stats::incMallocCount() {
	this->mallocCount++;
}

void Stats::incCallocCount() {
	this->callocCount++;
}

void Stats::incReallocCount() {
	this->reallocCount++;
}

void Stats::incFreeCount() {
	this->freeCount++;
}

void Stats::incInvalidReadCount() {
	this->invalidReadCount++;
}

void Stats::incReadFenceOverflow() {
	this->readFenceOverflow++;
}

void Stats::incReadFenceUnderflow() {
	this->readFenceUnderflow++;
}

void Stats::incInvalidWriteCount() {
	this->invalidWriteCount++;
}

void Stats::incWriteFenceOverflow() {
	this->writeFenceOverflow++;
}

void Stats::incWriteFenceUnderflow() {
	this->writeFenceUnderflow++;
}

void Stats::incInvalidFreeCount() {
	this->invalidFreeCount++;
}

void Stats::incMidFreeChunkCount() {
	this->midFreeChunkCount++;
}

void Stats::incFreeNullCount() {
	this->freeNullCount++;
}

void Stats::incFreeNotFoundCount() {
	this->freeNotFoundCount++;
}

void Stats::incInvalidReturnCount() {
	this->invalidReturnCount++;
}

/**
 * Display All information tracked by the pintool.
 */
void Stats::displayResults(MemList memlist, FILE *fp) {
	if(fp == NULL) {
		fp = stdin;
	}
	// Print out the initial object
	fprintf(fp, "==================================\n");
	fprintf(fp, "%-20s %s\n", "ACTIONS ", "COUNT");
	fprintf(fp, "----------------------------------\n");
	fprintf(fp, "%-20s %d\n", "allocations: ", this->allocCount);
	fprintf(fp, "%-20s %d\n", "    malloc", this->mallocCount);
	fprintf(fp, "%-20s %d\n", "    calloc", this->callocCount);
	fprintf(fp, "%-20s %d\n", "    realloc", this->reallocCount);
	fprintf(fp, "----------------------------------\n");
	fprintf(fp, "%-20s %d\n", "deallocations: ", this->freeCount);
	fprintf(fp, "%-20s %d\n", "    free ", this->freeCount);
	fprintf(fp, "----------------------------------\n");
	fprintf(fp, "%-20s %d\n", "invalid reads: ", this->invalidReadCount);
	fprintf(fp, "%-20s %d\n", "    underflow", this->readFenceUnderflow);
	fprintf(fp, "%-20s %d\n", "    overflow", this->readFenceOverflow);
	fprintf(fp, "----------------------------------\n");
	fprintf(fp, "%-20s %d\n", "invalid writes: ", this->invalidWriteCount);
	fprintf(fp, "%-20s %d\n", "    underflow", this->writeFenceUnderflow);
	fprintf(fp, "%-20s %d\n", "    overflow", this->writeFenceOverflow);
	fprintf(fp, "----------------------------------\n");
	fprintf(fp, "%-20s %d\n", "invalid frees: ", this->invalidFreeCount);
	fprintf(fp, "%-20s %d\n", "    not found", this->freeNotFoundCount);
	fprintf(fp, "%-20s %d\n", "    null free", this->freeNullCount);
	fprintf(fp, "%-20s %d\n", "    midchunk free", this->midFreeChunkCount);
	fprintf(fp, "----------------------------------\n");
	fprintf(fp, "%-20s %d\n", "invalid returns: ", this->invalidReturnCount);
	fprintf(fp, "==================================\n");
	// Print out the contents of memlist if it was provided.
	fprintf(fp, "%-20s\n", "MEMLIST CONTENTS");
	char buffer[2048];
	int lostMemory = 0;
	for(int i = 0; i < memlist.size(); i++) {
		MemoryAlloc alloc = memlist.get(i);
		fprintf(fp, "%s\n", alloc.toString(buffer, 2048));
		lostMemory += alloc.getTotalSize();
	}
	fprintf(fp, "----------------------------------\n");
	fprintf(fp, "%-20s %d\n", "list size: ", memlist.size());
	fprintf(fp, "%-20s %d %s\n", "total loss: ", lostMemory, "bytes");
	fprintf(fp, "==================================\n");
}
<commit_msg>Changed total size to reflect the user requested size instead of the size including the memoryfence.<commit_after>#include "stats.h"

Stats::Stats() {
	this->reset();
}

void Stats::reset() {
	/* Allocation Stats */
	this->allocCount = 0;
	this->mallocCount = 0;
	this->callocCount = 0;
	this->reallocCount = 0;
	/* Deallocation Stats */
	this->freeCount = 0;
	/* Invalid Free Stats*/
	this->invalidFreeCount = 0;
	this->midFreeChunkCount = 0;
	this->freeNullCount = 0;
	/* Invalid Read Stats */
	this->invalidReadCount = 0;
	this->readFenceOverflow = 0;
	this->readFenceUnderflow = 0;
	/* Invalid Write Stats */
	this->invalidWriteCount = 0;
	this->writeFenceOverflow = 0;
	this->writeFenceUnderflow = 0;
	/* Stack Smashing Stats */
	this->invalidReturnCount = 0;
}
/* Getters */
unsigned int Stats::getAllocCount() {
	return this->allocCount;
}

unsigned int Stats::getMallocCount() {
	return this->mallocCount;
}

unsigned int Stats::getCallocCount() {
	return this->callocCount;
}

unsigned int Stats::getReallocCount() {
	return this->reallocCount;
}

unsigned int Stats::getFreeCount() {
	return this->freeCount;
}
unsigned int Stats::getInvalidReadCount() {
	return this->invalidReadCount;
}

unsigned int Stats::getReadFenceOverflow() {
	return this->readFenceOverflow;
}

unsigned int Stats::getReadFenceUnderflow() {
	return this->readFenceUnderflow;
}

unsigned int Stats::getInvalidWriteCount() {
	return this->invalidWriteCount;
}

unsigned int Stats::getWriteFenceOverflow() {
	return this->writeFenceOverflow;
}

unsigned int Stats::getWriteFenceUnderflow() {
	return this->writeFenceUnderflow;
}

unsigned int Stats::getInvalidFreeCount() {
	return this->invalidFreeCount;
}

unsigned int Stats::getMidFreeChunkCount() {
	return this->midFreeChunkCount;
}

unsigned int Stats::getFreeNullCount() {
	return this->freeNullCount;
}

unsigned int Stats::getFreeNotFoundCount() {
	return this->freeNotFoundCount;
}

unsigned int Stats::getInvalidReturnCount() {
	return this->invalidReturnCount;
}

/* Setters */

void Stats::setAllocCount(unsigned int count) {
	this->allocCount = count;
}

void Stats::setMallocCount(unsigned int count) {
	this->allocCount = count;
}

void Stats::setCallocCount(unsigned int count) {
	this->allocCount = count;
}

void Stats::setReallocCount(unsigned int count) {
	this->allocCount = count;
}

void Stats::setFreeCount(unsigned int count) {
	this->freeCount = count;
}

void Stats::setInvalidReadCount(unsigned int count) {
	this->invalidReadCount = count;
}

void Stats::setReadFenceOverflow(unsigned int count) {
	this->readFenceOverflow = count;
}

void Stats::setReadFenceUnderflow(unsigned int count) {
	this->readFenceUnderflow = count;
}

void Stats::setInvalidWriteCount(unsigned int count) {
	this->invalidWriteCount = count;
}

void Stats::setWriteFenceOverflow(unsigned int count) {
	this->writeFenceOverflow = count;
}

void Stats::setWriteFenceUnderflow(unsigned int count) {
	this->writeFenceUnderflow = count;
}

void Stats::setInvalidFreeCount(unsigned int count) {
	this->invalidFreeCount = count;
}

void Stats::setMidFreeChunkCount(unsigned int count) {
	this->midFreeChunkCount = count;
}

void Stats::setFreeNullCount(unsigned int count) {
	this->freeNullCount = count;
}

void Stats::setInvalidReturnCount(unsigned int count) {
	this->invalidReturnCount = count;
}

void Stats::setFreeNotFoundCount(unsigned int count) {
	this->freeNotFoundCount = count;
}

/* Increment Count Functions */

void Stats::incAllocCount() {
	this->allocCount++;
}

void Stats::incMallocCount() {
	this->mallocCount++;
}

void Stats::incCallocCount() {
	this->callocCount++;
}

void Stats::incReallocCount() {
	this->reallocCount++;
}

void Stats::incFreeCount() {
	this->freeCount++;
}

void Stats::incInvalidReadCount() {
	this->invalidReadCount++;
}

void Stats::incReadFenceOverflow() {
	this->readFenceOverflow++;
}

void Stats::incReadFenceUnderflow() {
	this->readFenceUnderflow++;
}

void Stats::incInvalidWriteCount() {
	this->invalidWriteCount++;
}

void Stats::incWriteFenceOverflow() {
	this->writeFenceOverflow++;
}

void Stats::incWriteFenceUnderflow() {
	this->writeFenceUnderflow++;
}

void Stats::incInvalidFreeCount() {
	this->invalidFreeCount++;
}

void Stats::incMidFreeChunkCount() {
	this->midFreeChunkCount++;
}

void Stats::incFreeNullCount() {
	this->freeNullCount++;
}

void Stats::incFreeNotFoundCount() {
	this->freeNotFoundCount++;
}

void Stats::incInvalidReturnCount() {
	this->invalidReturnCount++;
}

/**
 * Display All information tracked by the pintool.
 */
void Stats::displayResults(MemList memlist, FILE *fp) {
	if(fp == NULL) {
		fp = stdin;
	}
	// Print out the initial object
	fprintf(fp, "==================================\n");
	fprintf(fp, "%-20s %s\n", "ACTIONS ", "COUNT");
	fprintf(fp, "----------------------------------\n");
	fprintf(fp, "%-20s %d\n", "allocations: ", this->allocCount);
	fprintf(fp, "%-20s %d\n", "    malloc", this->mallocCount);
	fprintf(fp, "%-20s %d\n", "    calloc", this->callocCount);
	fprintf(fp, "%-20s %d\n", "    realloc", this->reallocCount);
	fprintf(fp, "----------------------------------\n");
	fprintf(fp, "%-20s %d\n", "deallocations: ", this->freeCount);
	fprintf(fp, "%-20s %d\n", "    free ", this->freeCount);
	fprintf(fp, "----------------------------------\n");
	fprintf(fp, "%-20s %d\n", "invalid reads: ", this->invalidReadCount);
	fprintf(fp, "%-20s %d\n", "    underflow", this->readFenceUnderflow);
	fprintf(fp, "%-20s %d\n", "    overflow", this->readFenceOverflow);
	fprintf(fp, "----------------------------------\n");
	fprintf(fp, "%-20s %d\n", "invalid writes: ", this->invalidWriteCount);
	fprintf(fp, "%-20s %d\n", "    underflow", this->writeFenceUnderflow);
	fprintf(fp, "%-20s %d\n", "    overflow", this->writeFenceOverflow);
	fprintf(fp, "----------------------------------\n");
	fprintf(fp, "%-20s %d\n", "invalid frees: ", this->invalidFreeCount);
	fprintf(fp, "%-20s %d\n", "    not found", this->freeNotFoundCount);
	fprintf(fp, "%-20s %d\n", "    null free", this->freeNullCount);
	fprintf(fp, "%-20s %d\n", "    midchunk free", this->midFreeChunkCount);
	fprintf(fp, "----------------------------------\n");
	fprintf(fp, "%-20s %d\n", "invalid returns: ", this->invalidReturnCount);
	fprintf(fp, "==================================\n");
	// Print out the contents of memlist if it was provided.
	fprintf(fp, "%-20s\n", "MEMLIST CONTENTS");
	char buffer[2048];
	int lostMemory = 0;
	for(int i = 0; i < memlist.size(); i++) {
		MemoryAlloc alloc = memlist.get(i);
		fprintf(fp, "%s\n", alloc.toString(buffer, 2048));
		lostMemory += alloc.getUserSize();
	}
	fprintf(fp, "----------------------------------\n");
	fprintf(fp, "%-20s %d\n", "list size: ", memlist.size());
	fprintf(fp, "%-20s %d %s\n", "total loss: ", lostMemory, "bytes");
	fprintf(fp, "==================================\n");
}
<|endoftext|>


<commit_before>//
//  A* with strict expanded list
//
#pragma once

#include <functional>
#include <unordered_set>
#include "priority_queue.hpp"
#include "func.hpp"
#include "node.hpp"
#include "timing.hpp"

using std::size_t;
using std::string;
using std::move;
using std::max;
using std::any_of;
using std::unordered_set;

namespace search
{
    //
    //  Hfunc : heuristic function, (Node) -> h value
    //  ValidateFunc : unary predicate, check if a node is valid, (Node) -> bool
    //  Cfunc : cost function, (Node) -> c value
    //
    template<typename Hfunc, typename ValidateFunc, typename Cfunc = Cost<Node>>
    class AStarSEL
    {
    public:
        using Q = PriorityQueue<Node, Less<Node, Hfunc>>;
        using Expansions = unordered_set<State>;

        struct Get
        {
            Q q;
            size_t const& max_q_size;
            Expansions const& expansions;
            string const& final_path;
            long long const& run_time;
            bool const& is_found;
        } const last_run;

        AStarSEL()
            : last_run{ _q, _max_q_size, _expansions, _final_path, _run_time, _is_found }
        {
            reset();
        }

        auto operator()(State start, State goal, ValidateFunc validate) -> void
        {
            reset();
            Timing timing{ _run_time };
            search(start, goal, validate);
        }

    private:
        //
        //  Data members
        //
        Q _q;
        size_t _max_q_size;
        Expansions _expansions;
        string _final_path;
        long long _run_time;
        bool _is_found;
        //
        //  reset all data members
        //
        auto reset() -> void
        {
            _q.reset();
            _max_q_size = 0;
            _expansions.clear();
            _final_path.clear();
            _run_time = 0;
            _is_found = false;
        }
        //
        //  Core part
        //
        auto search(State start, State goal, ValidateFunc validate) -> void
        {
            for (_q.push({ "", start, goal }); !_q.empty() && _q.top().state() != goal; update_max_q_size())
            {
                auto curr = _q.pop();
                if (!has_expanded_into_state_of(curr))
                {
                    expand_into(curr.state());
                    for (auto const& child : curr.children(validate))
                        if (!has_expanded_into_state_of(child))
                            if (!_q.any(SameStateAs{ child.state() }))
                                _q.push(child);
                            else
                                _q.update_with_if(child, SameStateAs{ child.state() });
                }
            }

            if (_q.empty())
            {
                _final_path = "";
                _is_found = false;
            }
            else
            {
                _final_path = _q.top().path();
                _is_found = true;
            }
        }
        //
        //  helper for method search
        //
        auto update_max_q_size() -> void
        {
            _max_q_size = max(_max_q_size, _q.size());
        }
        //
        //  helper for method search
        //
        struct SameStateAs
        {
            State state_to_compare;
            auto operator()(Node const& n) const -> bool
            {
                return n.state() == state_to_compare;
            }
        };
        //
        //  helper for method search
        //
        auto has_expanded_into_state_of(Node const& n) const -> bool
        {
            auto const& e = _expansions;
            return any_of(e.cbegin(), e.cend(), [&](State s) {return s == n.state(); });
        }
        //
        //  helper for method search
        //
        auto expand_into(State s) -> void
        {
            _expansions.insert(s);
        }
    };
}<commit_msg>changed parameters for method search, not tested<commit_after>//
//  A* with strict expanded list
//
#pragma once

#include <functional>
#include <unordered_set>
#include "priority_queue.hpp"
#include "func.hpp"
#include "node.hpp"
#include "timing.hpp"

using std::size_t;
using std::string;
using std::move;
using std::max;
using std::any_of;
using std::unordered_set;

namespace search
{
    //
    //  Hfunc : heuristic function, (Node) -> h value
    //  ValidateFunc : unary predicate, check if a node is valid, (Node) -> bool
    //  Cfunc : cost function, (Node) -> c value
    //
    template<typename Hfunc, typename ValidateFunc, typename Cfunc = Cost<Node>>
    class AStarSEL
    {
    public:
        using Q = PriorityQueue<Node, Less<Node, Hfunc>>;
        using Expansions = unordered_set<State>;

        struct Get
        {
            Q q;
            size_t const& max_q_size;
            Expansions const& expansions;
            string const& final_path;
            long long const& run_time;
            bool const& is_found;
        } const last_run;

        AStarSEL()
            : last_run{ _q, _max_q_size, _expansions, _final_path, _run_time, _is_found }
        {
            reset();
        }

        auto operator()(State start, State goal, ValidateFunc validate) -> void
        {
            reset();
            Timing timing{ _run_time };
            search(start, goal, validate);
        }

    private:
        //
        //  Data members
        //
        Q _q;
        size_t _max_q_size;
        Expansions _expansions;
        string _final_path;
        long long _run_time;
        bool _is_found;
        //
        //  reset all data members
        //
        auto reset() -> void
        {
            _q.reset();
            _max_q_size = 0;
            _expansions.clear();
            _final_path.clear();
            _run_time = 0;
            _is_found = false;
        }
        //
        //  Core part
        //
        auto search(Node from, ValidateFunc validate) -> void
        {
            for (_q.push(from.state()); !_q.empty() && _q.top().state() != goal; update_max_q_size())
            {
                auto curr = _q.pop();
                if (!has_expanded_into_state_of(curr))
                {
                    expand_into(curr.state());
                    for (auto const& child : curr.children(validate))
                        if (!has_expanded_into_state_of(child))
                            if (!_q.any(SameStateAs{ child.state() }))
                                _q.push(child);
                            else
                                _q.update_with_if(child, SameStateAs{ child.state() });
                }
            }

            if (_q.empty())
            {
                _final_path = "";
                _is_found = false;
            }
            else
            {
                _final_path = _q.top().path();
                _is_found = true;
            }
        }
        //
        //  helper for method search
        //
        auto update_max_q_size() -> void
        {
            _max_q_size = max(_max_q_size, _q.size());
        }
        //
        //  helper for method search
        //
        struct SameStateAs
        {
            State state_to_compare;
            auto operator()(Node const& n) const -> bool
            {
                return n.state() == state_to_compare;
            }
        };
        //
        //  helper for method search
        //
        auto has_expanded_into_state_of(Node const& n) const -> bool
        {
            auto const& e = _expansions;
            return any_of(e.cbegin(), e.cend(), [&](State s) {return s == n.state(); });
        }
        //
        //  helper for method search
        //
        auto expand_into(State s) -> void
        {
            _expansions.insert(s);
        }
    };
}
<|endoftext|>


<commit_before>/* Copyright 2007-2015 QReal Research Group
 *
 * 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 "styledButton.h"

#include <QtWidgets/QLabel>
#include <QtWidgets/QStyle>

#include "brandManager/brandManager.h"
#include "circleWidget.h"

using namespace qReal;

StyledButton::StyledButton(const QString &text, const QString &icon
		, QBoxLayout::Direction direction, QWidget *parent)
	: QPushButton(parent)
{
	setMouseTracking(true);

	QBoxLayout * const layout = new QBoxLayout(direction);
	layout->setMargin(2);
	if (!icon.isEmpty()) {
		QWidget * const circleWidget = new CircleWidget(QSize(70, 70), icon);
		layout->addWidget(circleWidget);
		bindHighlightedOnHover(circleWidget);
	}

	QLabel * const textLabel = new QLabel(text);
	textLabel->setWordWrap(true);
	textLabel->setAttribute(Qt::WA_Hover);
	layout->addWidget(textLabel);
	bindHighlightedOnHover(textLabel);

	setFlat(true);
	setStyleSheet(BrandManager::styles()->startTabButtonStyle());
	setSizePolicy(QSizePolicy::Preferred, QSizePolicy::Preferred);
	setLayout(layout);
}

void StyledButton::bindHighlightedOnHover(QWidget * const widget)
{
	widget->setProperty("enabled", false);
	mChildren << widget;
}

void StyledButton::enterEvent(QEvent *event)
{
	QPushButton::enterEvent(event);
	highlight(true);
}

void StyledButton::leaveEvent(QEvent *event)
{
	QPushButton::leaveEvent(event);
	highlight(false);
}

void StyledButton::highlight(bool on)
{
	for (QWidget * const child : mChildren) {
		child->setProperty("enabled", on);
		child->style()->unpolish(child);
		child->style()->polish(child);
	}
}
<commit_msg>A dirty fix for styled button size, now 'Open interpreted diagram' on a start screen is not being cut.<commit_after>/* Copyright 2007-2015 QReal Research Group
 *
 * 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 "styledButton.h"

#include <QtWidgets/QLabel>
#include <QtWidgets/QStyle>

#include "brandManager/brandManager.h"
#include "circleWidget.h"

using namespace qReal;

StyledButton::StyledButton(const QString &text, const QString &icon
		, QBoxLayout::Direction direction, QWidget *parent)
	: QPushButton(parent)
{
	setMouseTracking(true);

	QBoxLayout * const layout = new QBoxLayout(direction);
	layout->setMargin(2);
	if (!icon.isEmpty()) {
		QWidget * const circleWidget = new CircleWidget(QSize(70, 70), icon);
		layout->addWidget(circleWidget);
		bindHighlightedOnHover(circleWidget);
	}

	QLabel * const textLabel = new QLabel(text);
	textLabel->setWordWrap(true);
	textLabel->setAttribute(Qt::WA_Hover);

	// Beginning from some version of Qt >= 5.5 QLabel does not take into account word wrap when calculating size hint,
	// so second line becomes clipped off. We use two lines for some buttons, so we manually set minimum height of a
	// label to fit two lines with our default font. Somebody who has internet connection shall provide more adequate
	// solution.
	textLabel->setMinimumHeight(55);

	layout->addWidget(textLabel);
	bindHighlightedOnHover(textLabel);

	setFlat(true);
	setStyleSheet(BrandManager::styles()->startTabButtonStyle());
	setSizePolicy(QSizePolicy::Preferred, QSizePolicy::Preferred);
	setLayout(layout);
}

void StyledButton::bindHighlightedOnHover(QWidget * const widget)
{
	widget->setProperty("enabled", false);
	mChildren << widget;
}

void StyledButton::enterEvent(QEvent *event)
{
	QPushButton::enterEvent(event);
	highlight(true);
}

void StyledButton::leaveEvent(QEvent *event)
{
	QPushButton::leaveEvent(event);
	highlight(false);
}

void StyledButton::highlight(bool on)
{
	for (QWidget * const child : mChildren) {
		child->setProperty("enabled", on);
		child->style()->unpolish(child);
		child->style()->polish(child);
	}
}
<|endoftext|>


<commit_before>#include "THMMesh.h"
#include "libmesh/node.h"
#include "libmesh/edge_edge2.h"
#include "libmesh/edge_edge3.h"
#include "libmesh/face_quad4.h"
#include "libmesh/face_quad9.h"

registerMooseObject("THMApp", THMMesh);

const BoundaryName THMMesh::INVALID_BOUNDARY_ID = "invalid_boundary_id";

template <>
InputParameters
validParams<THMMesh>()
{
  InputParameters params = validParams<MooseMesh>();
  // we do not allow renumbering, becuase we generate our meshes
  params.set<bool>("allow_renumbering") = false;
  return params;
}

THMMesh::THMMesh(const InputParameters & parameters)
  : MooseMesh(parameters),
    _dim(getParam<MooseEnum>("dim")),
    _next_node_id(0),
    _next_element_id(0),
    _next_subdomain_id(0),
    _next_boundary_id(0)
{
}

THMMesh::THMMesh(const THMMesh & other_mesh)
  : MooseMesh(other_mesh),
    _dim(other_mesh._dim),
    _next_node_id(other_mesh._next_node_id),
    _next_element_id(other_mesh._next_element_id),
    _next_subdomain_id(other_mesh._next_subdomain_id),
    _next_boundary_id(other_mesh._next_boundary_id)
{
}

unsigned int
THMMesh::dimension() const
{
  return _dim;
}

unsigned int
THMMesh::effectiveSpatialDimension() const
{
  return _dim;
}

MooseMesh &
THMMesh::clone() const
{
  mooseError("CRITICAL ERROR: calling clone() is not allowed and should not happen.");
}

std::unique_ptr<MooseMesh>
THMMesh::safeClone() const
{
  return libmesh_make_unique<THMMesh>(*this);
}

void
THMMesh::buildMesh()
{
  getMesh().set_spatial_dimension(_dim);
}

void
THMMesh::prep()
{
  prepare(true);
}

dof_id_type
THMMesh::getNextNodeId()
{
  unsigned int id = _next_node_id++;
  return id;
}

dof_id_type
THMMesh::getNextElementId()
{
  unsigned int id = _next_element_id++;
  return id;
}

Node *
THMMesh::addNode(const Point & pt)
{
  dof_id_type id = getNextNodeId();
  Node * node = _mesh->add_point(pt, id);
  node->set_unique_id() = id;
  return node;
}

Elem *
THMMesh::addElementEdge2(dof_id_type node0, dof_id_type node1)
{
  const auto pid = _mesh->comm().rank();
  dof_id_type elem_id = getNextElementId();

  Elem * elem = new Edge2;
  elem->set_id(elem_id);
  elem->processor_id() = pid;
  elem->set_unique_id() = elem_id;
  _mesh->add_elem(elem);
  elem->set_node(0) = _mesh->node_ptr(node0);
  elem->set_node(1) = _mesh->node_ptr(node1);
  return elem;
}

Elem *
THMMesh::addElementEdge3(dof_id_type node0, dof_id_type node1, dof_id_type node2)
{
  const auto pid = _mesh->comm().rank();
  dof_id_type elem_id = getNextElementId();

  Elem * elem = new Edge3;
  elem->set_id(elem_id);
  elem->processor_id() = pid;
  elem->set_unique_id() = elem_id;
  _mesh->add_elem(elem);
  elem->set_node(0) = _mesh->node_ptr(node0);
  elem->set_node(1) = _mesh->node_ptr(node1);
  elem->set_node(2) = _mesh->node_ptr(node2);
  return elem;
}

Elem *
THMMesh::addElementQuad4(dof_id_type node0, dof_id_type node1, dof_id_type node2, dof_id_type node3)
{
  const auto pid = _mesh->comm().rank();
  dof_id_type elem_id = getNextElementId();

  Elem * elem = new Quad4;
  elem->set_id(elem_id);
  elem->processor_id() = pid;
  elem->set_unique_id() = elem_id;
  _mesh->add_elem(elem);
  elem->set_node(0) = _mesh->node_ptr(node0);
  elem->set_node(1) = _mesh->node_ptr(node1);
  elem->set_node(2) = _mesh->node_ptr(node2);
  elem->set_node(3) = _mesh->node_ptr(node3);
  return elem;
}

Elem *
THMMesh::addElementQuad9(dof_id_type node0,
                         dof_id_type node1,
                         dof_id_type node2,
                         dof_id_type node3,
                         dof_id_type node4,
                         dof_id_type node5,
                         dof_id_type node6,
                         dof_id_type node7,
                         dof_id_type node8)
{
  const auto pid = _mesh->comm().rank();
  dof_id_type elem_id = getNextElementId();

  Elem * elem = new Quad9;
  elem->set_id(elem_id);
  elem->processor_id() = pid;
  elem->set_unique_id() = elem_id;
  _mesh->add_elem(elem);
  // vertices
  elem->set_node(0) = _mesh->node_ptr(node0);
  elem->set_node(1) = _mesh->node_ptr(node1);
  elem->set_node(2) = _mesh->node_ptr(node2);
  elem->set_node(3) = _mesh->node_ptr(node3);
  // mid-edges
  elem->set_node(4) = _mesh->node_ptr(node4);
  elem->set_node(5) = _mesh->node_ptr(node5);
  elem->set_node(6) = _mesh->node_ptr(node6);
  elem->set_node(7) = _mesh->node_ptr(node7);
  // center
  elem->set_node(8) = _mesh->node_ptr(node8);

  return elem;
}

SubdomainID
THMMesh::getNextSubdomainId()
{
  SubdomainID id = _next_subdomain_id++;
  return id;
}

BoundaryID
THMMesh::getNextBoundaryId()
{
  BoundaryID id = _next_boundary_id++;
  return id;
}
<commit_msg>Fixing incorrect type<commit_after>#include "THMMesh.h"
#include "libmesh/node.h"
#include "libmesh/edge_edge2.h"
#include "libmesh/edge_edge3.h"
#include "libmesh/face_quad4.h"
#include "libmesh/face_quad9.h"

registerMooseObject("THMApp", THMMesh);

const BoundaryName THMMesh::INVALID_BOUNDARY_ID = "invalid_boundary_id";

template <>
InputParameters
validParams<THMMesh>()
{
  InputParameters params = validParams<MooseMesh>();
  // we do not allow renumbering, becuase we generate our meshes
  params.set<bool>("allow_renumbering") = false;
  return params;
}

THMMesh::THMMesh(const InputParameters & parameters)
  : MooseMesh(parameters),
    _dim(getParam<MooseEnum>("dim")),
    _next_node_id(0),
    _next_element_id(0),
    _next_subdomain_id(0),
    _next_boundary_id(0)
{
}

THMMesh::THMMesh(const THMMesh & other_mesh)
  : MooseMesh(other_mesh),
    _dim(other_mesh._dim),
    _next_node_id(other_mesh._next_node_id),
    _next_element_id(other_mesh._next_element_id),
    _next_subdomain_id(other_mesh._next_subdomain_id),
    _next_boundary_id(other_mesh._next_boundary_id)
{
}

unsigned int
THMMesh::dimension() const
{
  return _dim;
}

unsigned int
THMMesh::effectiveSpatialDimension() const
{
  return _dim;
}

MooseMesh &
THMMesh::clone() const
{
  mooseError("CRITICAL ERROR: calling clone() is not allowed and should not happen.");
}

std::unique_ptr<MooseMesh>
THMMesh::safeClone() const
{
  return libmesh_make_unique<THMMesh>(*this);
}

void
THMMesh::buildMesh()
{
  getMesh().set_spatial_dimension(_dim);
}

void
THMMesh::prep()
{
  prepare(true);
}

dof_id_type
THMMesh::getNextNodeId()
{
  dof_id_type id = _next_node_id++;
  return id;
}

dof_id_type
THMMesh::getNextElementId()
{
  dof_id_type id = _next_element_id++;
  return id;
}

Node *
THMMesh::addNode(const Point & pt)
{
  dof_id_type id = getNextNodeId();
  Node * node = _mesh->add_point(pt, id);
  node->set_unique_id() = id;
  return node;
}

Elem *
THMMesh::addElementEdge2(dof_id_type node0, dof_id_type node1)
{
  const auto pid = _mesh->comm().rank();
  dof_id_type elem_id = getNextElementId();

  Elem * elem = new Edge2;
  elem->set_id(elem_id);
  elem->processor_id() = pid;
  elem->set_unique_id() = elem_id;
  _mesh->add_elem(elem);
  elem->set_node(0) = _mesh->node_ptr(node0);
  elem->set_node(1) = _mesh->node_ptr(node1);
  return elem;
}

Elem *
THMMesh::addElementEdge3(dof_id_type node0, dof_id_type node1, dof_id_type node2)
{
  const auto pid = _mesh->comm().rank();
  dof_id_type elem_id = getNextElementId();

  Elem * elem = new Edge3;
  elem->set_id(elem_id);
  elem->processor_id() = pid;
  elem->set_unique_id() = elem_id;
  _mesh->add_elem(elem);
  elem->set_node(0) = _mesh->node_ptr(node0);
  elem->set_node(1) = _mesh->node_ptr(node1);
  elem->set_node(2) = _mesh->node_ptr(node2);
  return elem;
}

Elem *
THMMesh::addElementQuad4(dof_id_type node0, dof_id_type node1, dof_id_type node2, dof_id_type node3)
{
  const auto pid = _mesh->comm().rank();
  dof_id_type elem_id = getNextElementId();

  Elem * elem = new Quad4;
  elem->set_id(elem_id);
  elem->processor_id() = pid;
  elem->set_unique_id() = elem_id;
  _mesh->add_elem(elem);
  elem->set_node(0) = _mesh->node_ptr(node0);
  elem->set_node(1) = _mesh->node_ptr(node1);
  elem->set_node(2) = _mesh->node_ptr(node2);
  elem->set_node(3) = _mesh->node_ptr(node3);
  return elem;
}

Elem *
THMMesh::addElementQuad9(dof_id_type node0,
                         dof_id_type node1,
                         dof_id_type node2,
                         dof_id_type node3,
                         dof_id_type node4,
                         dof_id_type node5,
                         dof_id_type node6,
                         dof_id_type node7,
                         dof_id_type node8)
{
  const auto pid = _mesh->comm().rank();
  dof_id_type elem_id = getNextElementId();

  Elem * elem = new Quad9;
  elem->set_id(elem_id);
  elem->processor_id() = pid;
  elem->set_unique_id() = elem_id;
  _mesh->add_elem(elem);
  // vertices
  elem->set_node(0) = _mesh->node_ptr(node0);
  elem->set_node(1) = _mesh->node_ptr(node1);
  elem->set_node(2) = _mesh->node_ptr(node2);
  elem->set_node(3) = _mesh->node_ptr(node3);
  // mid-edges
  elem->set_node(4) = _mesh->node_ptr(node4);
  elem->set_node(5) = _mesh->node_ptr(node5);
  elem->set_node(6) = _mesh->node_ptr(node6);
  elem->set_node(7) = _mesh->node_ptr(node7);
  // center
  elem->set_node(8) = _mesh->node_ptr(node8);

  return elem;
}

SubdomainID
THMMesh::getNextSubdomainId()
{
  SubdomainID id = _next_subdomain_id++;
  return id;
}

BoundaryID
THMMesh::getNextBoundaryId()
{
  BoundaryID id = _next_boundary_id++;
  return id;
}
<|endoftext|>


<commit_before>#include <THtml.h>
#include <TROOT.h>
#include <TSystem.h>

void mkhtml (char *macro=0, Int_t force=0) {
// to run this macro, you must have the correct .rootrc file
// in your galice directory.
// The gAlice classes summary documentation go to directory html
// The gAlice classes source  documentation go to directory html/src
// The example macros documentation go to directory html/examples

  // gROOT->LoadMacro("loadlibs.C");
  // loadlibs();
  THtml html;
  html.SetProductName("AliRoot");
  if(macro) {
    gROOT->LoadMacro(macro);
    html.Convert(macro,"Example Macro");
  } else {
    gSystem->Load("liblhapdf.so");      // Parton density functions
    gSystem->Load("libEGPythia6.so");   // TGenerator interface
    gSystem->Load("libpythia6.so");     // Pythia
    gSystem->Load("libAliPythia6.so");  // ALICE specific implementations

    gSystem->Load("libANALYSIS.so");
    gSystem->Load("libANALYSISalice.so");
    gSystem->Load("libANALYSIScalib.so");
    gSystem->Load("libEventMixing.so");

    gSystem->Load("libPWG0base.so");
    gSystem->Load("libPWG0dep.so");
    gSystem->Load("libPWG0selectors.so");

    gSystem->Load("libTENDER.so");
    gSystem->Load("libPWGPP.so");

    gSystem->Load("libCORRFW.so");
    gSystem->Load("libPWG2AOD.so");
    gSystem->Load("libPWG2ebye.so");
    gSystem->Load("libPWG2evchar.so");
    gSystem->Load("libPWG2femtoscopy.so");
    gSystem->Load("libPWG2femtoscopyUser.so");
    gSystem->Load("libPWG2flowCommon.so");
    gSystem->Load("libPWG2flowTasks.so");
    gSystem->Load("libPWG2forward.so");
    gSystem->Load("libPWG2kink.so");
    gSystem->Load("libPWG2resonances.so");
    gSystem->Load("libPWG2spectra.so");
    gSystem->Load("libPWG2unicor.so");

    gSystem->Load("libPWGHFbase.so");
    gSystem->Load("libPWGHFhfe.so");
    gSystem->Load("libPWGmuondep.so");
    gSystem->Load("libPWGmuon.so");
    gSystem->Load("libPWGHFvertexingHF.so");

    gSystem->Load("libJETAN.so");
    gSystem->Load("libPWG4CaloCalib.so");
    gSystem->Load("libPWG4GammaConv.so");
    gSystem->Load("libPWG4JetTasks.so");
    gSystem->Load("libPWG4omega3pi.so");
    gSystem->Load("libPWG4PartCorrBase.so");
    gSystem->Load("libPWG4PartCorrDep.so");

    // EVE
    gSystem->Load("libEve.so");
    gSystem->Load("libEveBase.so");
    gSystem->Load("libEveDet.so");
    gSystem->Load("libEveHLT.so");

    html.MakeAll(force,"[A-Z]*");
  }
}
<commit_msg>Updated list of libraries<commit_after>#include <THtml.h>
#include <TROOT.h>
#include <TSystem.h>

void mkhtml (char *macro=0, Int_t force=0) {
// to run this macro, you must have the correct .rootrc file
// in your galice directory.
// The gAlice classes summary documentation go to directory html
// The gAlice classes source  documentation go to directory html/src
// The example macros documentation go to directory html/examples

  // gROOT->LoadMacro("loadlibs.C");
  // loadlibs();
  THtml html;
  html.SetProductName("AliRoot");
  if(macro) {
    gROOT->LoadMacro(macro);
    html.Convert(macro,"Example Macro");
  } else {
    gSystem->Load("liblhapdf.so");      // Parton density functions
    gSystem->Load("libEGPythia6.so");   // TGenerator interface
    gSystem->Load("libpythia6.so");     // Pythia
    gSystem->Load("libAliPythia6.so");  // ALICE specific implementations

    // ANALYSIS
    gSystem->Load("libANALYSIS.so");
    gSystem->Load("libANALYSISalice.so");
    gSystem->Load("libANALYSIScalib.so");
    gSystem->Load("libEventMixing.so");
    gSystem->Load("libTENDER.so");
    gSystem->Load("libTENDERSupplies.so");

    // CORRFW
    gSystem->Load("libCORRFW.so"); //

    // PWG
    gSystem->Load("libPWGTools.so");
    gSystem->Load("libPWGGlauber.so");
    gSystem->Load("libPWGflowBase.so");
    gSystem->Load("libPWGflowTasks.so");
    gSystem->Load("libPWGmuon.so");
    gSystem->Load("libPWGmuondep.so");
    gSystem->Load("libPWGEMCAL.so");
    gSystem->Load("libPWGCaloTrackCorrBase.so");

    // PWGCF
    gSystem->Load("libPWGCFCorrelationsBase.so");
    gSystem->Load("libPWGCFCorrelationsDPhi.so");
    gSystem->Load("libPWGCFCorrelationsJCORRAN.so");
    gSystem->Load("libPWGCFChaoticity.so");
    gSystem->Load("libPWGCFFEMTOSCOPYAOD.so");
    gSystem->Load("libPWGCFfemtoscopy.so");
    gSystem->Load("libPWGCFfemtoscopyUser.so");
    gSystem->Load("libPWGCFunicor.so");
    gSystem->Load("libPWGCFebye.so");
    //PH    gSystem->Load("libPWGCFK0Analysis.so");

    // PWGDQ
    gSystem->Load("libPWGDQbase.so");
    gSystem->Load("libPWGDQdielectron.so");

    // PWGGA
    gSystem->Load("libPWGGACaloTasks.so");
    gSystem->Load("libPWGGACaloTrackCorrelations.so");
    gSystem->Load("libPWGGAEMCALTasks.so");
    gSystem->Load("libPWGGAGammaConv.so");
    gSystem->Load("libPWGGAPHOSTasks.so");

    // PWGHF
    gSystem->Load("libPWGHFbase.so");
    gSystem->Load("libPWGHFhfe.so");
    gSystem->Load("libPWGHFcorrelationHF.so");
    gSystem->Load("libPWGHFvertexingHF.so");

    // PWGJE
    gSystem->Load("libJETAN.so");
    // gSystem->Load("libPWGJE.so");
    // gSystem->Load("libPWGJEEMCALJetTasks.so");

    // PWGLF
    gSystem->Load("libPWGLFSTRANGENESS.so");
    gSystem->Load("libPWGLFforward.so");
    gSystem->Load("libPWGLFforward2.so");
    gSystem->Load("libPWGLFresonances.so");
    gSystem->Load("libPWGLFrsnextra.so");
    gSystem->Load("libPWGLFspectra.so");
    // gSystem->Load("libPWGLFtotEt.so")

    // PWGPP
    gSystem->Load("libPWGPP.so");
    gSystem->Load("libPWGPPMUONdep.so");
    gSystem->Load("libPWGPPMUONlite.so");
    gSystem->Load("libPWGPPevchar.so");

    // PWGUD

    gSystem->Load("libPWGUDbase.so");
    gSystem->Load("libPWGUDFP.so");
    gSystem->Load("libPWGUDdiffractive.so");
    gSystem->Load("libPWGUDselectors.so");

    // EVE
    gSystem->Load("libEve.so");
    gSystem->Load("libEveBase.so");
    gSystem->Load("libEveDet.so");
    gSystem->Load("libEveHLT.so");

    html.MakeAll(force,"[A-Z]*");
  }
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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 <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <vector>

#include <ImfArray.h>
#include <ImfInputPart.h>
#include <ImfMultiPartInputFile.h>
#include <ImfChannelList.h>
#include <ImfRgbaFile.h>
#include <ImfStdIO.h>

using namespace OPENEXR_IMF_INTERNAL_NAMESPACE;
using IMATH_NAMESPACE::Box2i;
using std::vector;
using std::max;
namespace {

static void readSingle(IStream& is) {
  try {

    RgbaInputFile in(is);

    const Box2i &dw = in.dataWindow();

    int w = dw.max.x - dw.min.x + 1;
    int dx = dw.min.x;

    if (w > (1 << 24)) return;

    Array<Rgba> pixels(w);
    in.setFrameBuffer(&pixels[-dx], 1, 0);

    // read up to 10,000 scanlines from file
    int step = max( 1 , (dw.max.y - dw.min.y + 1) / 10000);
    for (int y = dw.min.y; y <= dw.max.y; y+=step) in.readPixels(y);

  } catch (...) {
  }
}

static void readMulti(IStream& is) {
  MultiPartInputFile *file;
  try {
    file = new MultiPartInputFile(is);
  } catch (...) {
    return;
  }

  for (int p = 0; p < file->parts(); p++) {
    InputPart *in;
    try {
      in = new InputPart(*file, p);
    } catch (...) {
      continue;
    }

    try {
      const Box2i &dw = in->header().dataWindow();

      int w = dw.max.x - dw.min.x + 1;
      int dx = dw.min.x;

      if (w > (1 << 24))
      {
	      throw std::logic_error("ignoring - very wide datawindow\n");
      }

      FrameBuffer i;
      //
      // read all channels present (later channels will overwrite earlier ones)
      vector<half> otherChannels(w);
      const ChannelList& channelList = in->header().channels();
      for (ChannelList::ConstIterator c = channelList.begin() ; c != channelList.end() ; ++c )
      {
	      i.insert(c.name(),Slice(HALF, (char*)&otherChannels[0] , sizeof(half) , 0 ));
      }

      // always try to read RGBA even if not present in file
      Array<Rgba> pixels(w);
      i.insert("R", Slice(HALF, (char *)&(pixels[-dx].r), sizeof(Rgba), 0));
      i.insert("G", Slice(HALF, (char *)&(pixels[-dx].g), sizeof(Rgba), 0));
      i.insert("B", Slice(HALF, (char *)&(pixels[-dx].b), sizeof(Rgba), 0));
      i.insert("A", Slice(HALF, (char *)&(pixels[-dx].a), sizeof(Rgba), 0));

      in->setFrameBuffer(i);
      // read up to 10,000 scanlines from file
      int step = max( 1 , (dw.max.y - dw.min.y + 1) / 10000);
      for (int y = dw.min.y; y <= dw.max.y; y+=step) in->readPixels(y);
    } catch (...) {
    }

    delete in;
  }

  delete file;
}

}  // namespace

extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
  const std::string s(reinterpret_cast<const char*>(data), size);
  {
    StdISStream is;
    is.str(s);
    readSingle(is);
  }
  {
    StdISStream is;
    is.str(s);
    readMulti(is);
  }
  return 0;
}
<commit_msg>[openexr] fix reading extra EXR channels in readMulti (#4461)<commit_after>// Copyright 2020 Google LLC
//
// 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 <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <vector>

#include <ImfArray.h>
#include <ImfInputPart.h>
#include <ImfMultiPartInputFile.h>
#include <ImfChannelList.h>
#include <ImfRgbaFile.h>
#include <ImfStdIO.h>

using namespace OPENEXR_IMF_INTERNAL_NAMESPACE;
using IMATH_NAMESPACE::Box2i;
using std::vector;
using std::max;
namespace {

static void readSingle(IStream& is) {
  try {

    RgbaInputFile in(is);

    const Box2i &dw = in.dataWindow();

    int w = dw.max.x - dw.min.x + 1;
    int dx = dw.min.x;

    if (w > (1 << 24)) return;

    Array<Rgba> pixels(w);
    in.setFrameBuffer(&pixels[-dx], 1, 0);

    // read up to 10,000 scanlines from file
    int step = max( 1 , (dw.max.y - dw.min.y + 1) / 10000);
    for (int y = dw.min.y; y <= dw.max.y; y+=step) in.readPixels(y);

  } catch (...) {
  }
}

static void readMulti(IStream& is) {
  MultiPartInputFile *file;
  try {
    file = new MultiPartInputFile(is);
  } catch (...) {
    return;
  }

  for (int p = 0; p < file->parts(); p++) {
    InputPart *in;
    try {
      in = new InputPart(*file, p);
    } catch (...) {
      continue;
    }

    try {
      const Box2i &dw = in->header().dataWindow();

      int w = dw.max.x - dw.min.x + 1;
      int dx = dw.min.x;

      if (w > (1 << 24))
      {
	      throw std::logic_error("ignoring - very wide datawindow\n");
      }

      FrameBuffer i;
      //
      // read all channels present (later channels will overwrite earlier ones)
      vector<half> otherChannels(w);
      const ChannelList& channelList = in->header().channels();
      for (ChannelList::ConstIterator c = channelList.begin() ; c != channelList.end() ; ++c )
      {
	      i.insert(c.name(),Slice(HALF, (char*)&otherChannels[-dx] , sizeof(half) , 0 ));
      }

      // always try to read RGBA even if not present in file
      Array<Rgba> pixels(w);
      i.insert("R", Slice(HALF, (char *)&(pixels[-dx].r), sizeof(Rgba), 0));
      i.insert("G", Slice(HALF, (char *)&(pixels[-dx].g), sizeof(Rgba), 0));
      i.insert("B", Slice(HALF, (char *)&(pixels[-dx].b), sizeof(Rgba), 0));
      i.insert("A", Slice(HALF, (char *)&(pixels[-dx].a), sizeof(Rgba), 0));

      in->setFrameBuffer(i);
      // read up to 10,000 scanlines from file
      int step = max( 1 , (dw.max.y - dw.min.y + 1) / 10000);
      for (int y = dw.min.y; y <= dw.max.y; y+=step) in->readPixels(y);
    } catch (...) {
    }

    delete in;
  }

  delete file;
}

}  // namespace

extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
  const std::string s(reinterpret_cast<const char*>(data), size);
  {
    StdISStream is;
    is.str(s);
    readSingle(is);
  }
  {
    StdISStream is;
    is.str(s);
    readMulti(is);
  }
  return 0;
}
<|endoftext|>


<commit_before>/* -*- Mode: C++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */
/*************************************************************************
 *
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * Copyright 2000, 2010 Oracle and/or its affiliates.
 *
 * OpenOffice.org - a multi-platform office productivity suite
 *
 * This file is part of OpenOffice.org.
 *
 * OpenOffice.org is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License version 3
 * only, as published by the Free Software Foundation.
 *
 * OpenOffice.org is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License version 3 for more details
 * (a copy is included in the LICENSE file that accompanied this code).
 *
 * You should have received a copy of the GNU Lesser General Public License
 * version 3 along with OpenOffice.org.  If not, see
 * <http://www.openoffice.org/license.html>
 * for a copy of the LGPLv3 License.
 *
 ************************************************************************/

#ifndef _OWNLIST_HXX
#define _OWNLIST_HXX

#include "svl/svldllapi.h"
#include <tools/stream.hxx>
#include <com/sun/star/uno/Sequence.hxx>
#include <vector>

namespace com { namespace sun { namespace star {
    namespace beans {
        struct PropertyValue;
    }
}}}

//=========================================================================
class SvCommand
/*  [Beschreibung]

    Enth"alt einen String, welcher das Kommando angibt und eine weiteren
    String, der das Argument des Kommandos bildet. W"urde solch ein
    Kommando "uber die Kommandozeile angegeben werden, s"ahe es wie folgt
    aus: Kommando = Argument.
*/
{
    String  aCommand;
    String  aArgument;
public:
                    SvCommand() {}
                    SvCommand( const String & rCommand, const String & rArg )
                    {
                        aCommand = rCommand;
                        aArgument = rArg;
                    }
    const String & GetCommand() const { return aCommand; }
    const String & GetArgument() const { return aArgument; }

    friend SvStream& operator >> ( SvStream& rStm, SvCommand & rThis )
                    {
                        rStm.ReadByteString( rThis.aCommand, gsl_getSystemTextEncoding() );
                        rStm.ReadByteString( rThis.aArgument, gsl_getSystemTextEncoding() );
                        return rStm;
                    }
    friend SvStream& operator << ( SvStream& rStm, const SvCommand & rThis )
                    {
                        rStm.WriteByteString( rThis.aCommand, gsl_getSystemTextEncoding() );
                        rStm.WriteByteString( rThis.aArgument, gsl_getSystemTextEncoding() );
                        return rStm;
                    }
};

typedef ::std::vector< SvCommand > SvCommandList_impl;

//=========================================================================
class SVL_DLLPUBLIC SvCommandList
/*  [Beschreibung]

    Die Liste enth"alt Objekte vom Typ SvCommand. Wird ein Objekt
    eingef"ugt, dann wird es kopiert und das neue Objekt wird
    in die Liste gestellt.
*/
{
private:
    SvCommandList_impl  aCommandList;

public:
    SvCommand&      Append( const String & rCommand, const String & rArg );
    sal_Bool        AppendCommands( const String & rCmd, sal_uInt16 * pEaten );
    String          GetCommands() const;

    sal_Bool FillFromSequence( const com::sun::star::uno::Sequence < com::sun::star::beans::PropertyValue >& );
    void FillSequence( com::sun::star::uno::Sequence < com::sun::star::beans::PropertyValue >& );

    size_t          size() const { return aCommandList.size(); }

    SvCommand       operator[]( size_t i) {
                        return aCommandList[ i ];
                    }

    void            clear() {
                        aCommandList.clear();
                    }
};

#endif // _OWNLIST_HXX

/* vim:set shiftwidth=4 softtabstop=4 expandtab: */
<commit_msg>ditch other unused stream operators<commit_after>/* -*- Mode: C++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */
/*************************************************************************
 *
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * Copyright 2000, 2010 Oracle and/or its affiliates.
 *
 * OpenOffice.org - a multi-platform office productivity suite
 *
 * This file is part of OpenOffice.org.
 *
 * OpenOffice.org is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License version 3
 * only, as published by the Free Software Foundation.
 *
 * OpenOffice.org is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License version 3 for more details
 * (a copy is included in the LICENSE file that accompanied this code).
 *
 * You should have received a copy of the GNU Lesser General Public License
 * version 3 along with OpenOffice.org.  If not, see
 * <http://www.openoffice.org/license.html>
 * for a copy of the LGPLv3 License.
 *
 ************************************************************************/

#ifndef _OWNLIST_HXX
#define _OWNLIST_HXX

#include "svl/svldllapi.h"
#include <tools/stream.hxx>
#include <com/sun/star/uno/Sequence.hxx>
#include <vector>

namespace com { namespace sun { namespace star {
    namespace beans {
        struct PropertyValue;
    }
}}}

//=========================================================================
class SvCommand
/*  [Beschreibung]

    Enth"alt einen String, welcher das Kommando angibt und eine weiteren
    String, der das Argument des Kommandos bildet. W"urde solch ein
    Kommando "uber die Kommandozeile angegeben werden, s"ahe es wie folgt
    aus: Kommando = Argument.
*/
{
    String  aCommand;
    String  aArgument;
public:
                    SvCommand() {}
                    SvCommand( const String & rCommand, const String & rArg )
                    {
                        aCommand = rCommand;
                        aArgument = rArg;
                    }
    const String & GetCommand() const { return aCommand; }
    const String & GetArgument() const { return aArgument; }
};

typedef ::std::vector< SvCommand > SvCommandList_impl;

//=========================================================================
class SVL_DLLPUBLIC SvCommandList
/*  [Beschreibung]

    Die Liste enth"alt Objekte vom Typ SvCommand. Wird ein Objekt
    eingef"ugt, dann wird es kopiert und das neue Objekt wird
    in die Liste gestellt.
*/
{
private:
    SvCommandList_impl  aCommandList;

public:
    SvCommand&      Append( const String & rCommand, const String & rArg );
    sal_Bool        AppendCommands( const String & rCmd, sal_uInt16 * pEaten );
    String          GetCommands() const;

    sal_Bool FillFromSequence( const com::sun::star::uno::Sequence < com::sun::star::beans::PropertyValue >& );
    void FillSequence( com::sun::star::uno::Sequence < com::sun::star::beans::PropertyValue >& );

    size_t          size() const { return aCommandList.size(); }

    SvCommand       operator[]( size_t i) {
                        return aCommandList[ i ];
                    }

    void            clear() {
                        aCommandList.clear();
                    }
};

#endif // _OWNLIST_HXX

/* vim:set shiftwidth=4 softtabstop=4 expandtab: */
<|endoftext|>


<commit_before>#include "threadPoolAsync.h"
#include "shared/defer.tcc"

#include <thread>

#include "../simpleServer/mt.h"

#include "exceptions.h"

namespace simpleServer {
using ondra_shared::defer;

ThreadPoolAsyncImpl::~ThreadPoolAsyncImpl() {
	stop();
}

void ThreadPoolAsyncImpl::cancel(const AsyncResource& resource) {
	Sync _(lock);
	auto sz = cQueue.size();
	while (sz) {
		PStreamEventDispatcher d = cQueue.front();
		cQueue.pop();
		d->cancel(resource);
		cQueue.push(d);
		--sz;
	}
}

void ThreadPoolAsyncImpl::checkThreadCount() {
	while (static_cast<unsigned int>(threadCount.getCounter()) < reqThreadCount) {
		RefCntPtr<ThreadPoolAsyncImpl> me(this);
		runThread([me] {
			me->worker();
		});
		threadCount.inc();
	}
}

class ThreadPoolAsyncImpl::InvokeNetworkDispatcher {
public:
	InvokeNetworkDispatcher(ThreadPoolAsyncImpl &owner, const PStreamEventDispatcher &sed)
		:owner(owner)
		,sed(sed) {}
	void operator()() const {
/*		auto t = sed->wait();
		if (t == nullptr) {
			disp.quit();
		} else {
			disp.dispatch(*this);
		}
		t();*/
		owner.waitForTask(sed);
	}
protected:
	ThreadPoolAsyncImpl &owner;
	PStreamEventDispatcher sed;
};

void ThreadPoolAsyncImpl::onExitDispatcher(const PStreamEventDispatcher &sed) noexcept {
	Sync _(lock);
	auto ccnt = cQueue.size();
	for (decltype(ccnt) i = 0; i < ccnt; i++) {
		const PStreamEventDispatcher &x = cQueue.front();
		if (x == sed) {
			cQueue.pop();
			if (cQueue.empty()) {
				dQueue.quit();
			}
			break;
		} else {
			cQueue.push(x);
			cQueue.pop();
		}
	}

}

void ThreadPoolAsyncImpl::waitForTask(const PStreamEventDispatcher &sed) noexcept {
			auto t = sed->wait();
			if (t == nullptr) {
				onExitDispatcher(sed);
				return;
			} else {
				dQueue.dispatch(InvokeNetworkDispatcher(*this, sed));
			}
			t();
}

PStreamEventDispatcher ThreadPoolAsyncImpl::getListener() {
	Sync _(lock);
	PStreamEventDispatcher lst,lst2;
	while (cQueue.size() < reqDispatcherCount) {
		lst = AbstractStreamEventDispatcher::create();
		cQueue.push(lst);
		dQueue.dispatch(InvokeNetworkDispatcher(*this, lst));
	}
	while (cQueue.size() > reqDispatcherCount) {
		//pop extra dispatcher
		//they are still in queue, but eventually disappear because they no longer receive a work
		cQueue.pop();
	}
	checkThreadCount();
	lst = cQueue.front();
	cQueue.pop();
	cQueue.push(lst);
	return lst;
}

void ThreadPoolAsyncImpl::runAsync(const AsyncResource& resource, int timeout, const CompletionFn &fn) {

	if (exitFlag) {
		defer >> std::bind(fn, asyncCancel);
		return;
	}

	unsigned int tries = 0;
	auto retry = [&] {
		Sync _(lock);
		++tries;
		if (tries >= cQueue.size())  {
			reqDispatcherCount++;
			if (reqDispatcherCount>reqThreadCount) reqThreadCount++;
		}
	};
	do {
		try {
			auto lst = getListener();
			if (lst->getPendingCount()>=taskLimit) {
				retry();
			} else {
				lst->runAsync(resource,timeout, fn);
				return;
			}
		} catch (OutOfSpaceException) {
			retry();
		}
	}while(true);
}

void ThreadPoolAsyncImpl::setCountOfDispatchers(unsigned int count) {
	Sync _(lock);
	reqDispatcherCount = count;
}

void ThreadPoolAsyncImpl::setCountOfThreads(unsigned int count) {
	Sync _(lock);
	reqThreadCount = count;
}

void ThreadPoolAsyncImpl::stop() {

	{
		Sync _(lock);
		exitFlag = true;
		auto sz = cQueue.size();
		while (sz) {
			PStreamEventDispatcher d = cQueue.front();
			cQueue.pop();
			cQueue.push(d);
			d->stop();
			--sz;
		}
	}
	threadCount.wait();

}

void ThreadPoolAsyncImpl::setTasksPerDispLimit(unsigned int count) {
	Sync _(lock);
	taskLimit = count;

}

void ThreadPoolAsyncImpl::worker() noexcept {

	using namespace ondra_shared;

	DeferContext defer(ondra_shared::defer_root);

	for(;;) {

		if (!dQueue.pump()) {
			dQueue.quit();
			defer_yield();
			threadCount.dec();
			return;
		}

		defer_yield();
/*		PStreamEventDispatcher lst = tQueue.pop();
		if (lst == nullptr) {
			tQueue.push(nullptr);
			threadCount.dec();
			return;
		}

		auto t = lst->wait();
		tQueue.push(lst);
		if (t == nullptr) {
			threadCount.dec();
			return;
		}

		t();
*/
		if (static_cast<unsigned int>(threadCount.getCounter()) > reqThreadCount) {
			Sync _(lock);
			if (static_cast<unsigned int>(threadCount.getCounter()) > reqThreadCount) {
				threadCount.dec();
				return;
			}
		}
	}
}

AsyncProvider ThreadPoolAsync::create(unsigned int numThreads, unsigned int numListeners, unsigned int tasksPerDispLimit) {
	ThreadPoolAsync provider (new ThreadPoolAsyncImpl);
	provider.setCountOfDispatchers(numListeners);
	provider.setCountOfThreads(numThreads);
	provider.setTasksPerDispLimit(tasksPerDispLimit);
	return provider;
}

void ThreadPoolAsync::setCountOfDispatchers(unsigned int count) {
	ptr->setCountOfDispatchers(count);
}

void ThreadPoolAsync::setCountOfThreads(unsigned int count) {
	ptr->setCountOfThreads(count);
}

void ThreadPoolAsync::stop() {
	ptr->stop();
}

ThreadPoolAsync::~ThreadPoolAsync() {
	if (!ptr->isShared()) {
		ptr->stop();
	}
}

void ThreadPoolAsyncImpl::runAsync(const CustomFn& completion) {

	if (exitFlag) {
		defer >> completion;
		return;
	}

	if (reqThreadCount > reqDispatcherCount) {
		checkThreadCount();
		dQueue.dispatch(completion);
	} else {
		auto lst = getListener();
		lst->runAsync(completion);
	}
}

void ThreadPoolAsync::setTasksPerDispLimit(unsigned int count) {
	ptr->setTasksPerDispLimit(count);
}

ThreadPoolAsync::operator AsyncProvider() const {
	return AsyncProvider(ptr);
}

}

<commit_msg>fix warning  error: catching polymorphic type ‘class simpleServer::OutOfSpaceException’ by value<commit_after>#include "threadPoolAsync.h"
#include "shared/defer.tcc"

#include <thread>

#include "../simpleServer/mt.h"

#include "exceptions.h"

namespace simpleServer {
using ondra_shared::defer;

ThreadPoolAsyncImpl::~ThreadPoolAsyncImpl() {
	stop();
}

void ThreadPoolAsyncImpl::cancel(const AsyncResource& resource) {
	Sync _(lock);
	auto sz = cQueue.size();
	while (sz) {
		PStreamEventDispatcher d = cQueue.front();
		cQueue.pop();
		d->cancel(resource);
		cQueue.push(d);
		--sz;
	}
}

void ThreadPoolAsyncImpl::checkThreadCount() {
	while (static_cast<unsigned int>(threadCount.getCounter()) < reqThreadCount) {
		RefCntPtr<ThreadPoolAsyncImpl> me(this);
		runThread([me] {
			me->worker();
		});
		threadCount.inc();
	}
}

class ThreadPoolAsyncImpl::InvokeNetworkDispatcher {
public:
	InvokeNetworkDispatcher(ThreadPoolAsyncImpl &owner, const PStreamEventDispatcher &sed)
		:owner(owner)
		,sed(sed) {}
	void operator()() const {
/*		auto t = sed->wait();
		if (t == nullptr) {
			disp.quit();
		} else {
			disp.dispatch(*this);
		}
		t();*/
		owner.waitForTask(sed);
	}
protected:
	ThreadPoolAsyncImpl &owner;
	PStreamEventDispatcher sed;
};

void ThreadPoolAsyncImpl::onExitDispatcher(const PStreamEventDispatcher &sed) noexcept {
	Sync _(lock);
	auto ccnt = cQueue.size();
	for (decltype(ccnt) i = 0; i < ccnt; i++) {
		const PStreamEventDispatcher &x = cQueue.front();
		if (x == sed) {
			cQueue.pop();
			if (cQueue.empty()) {
				dQueue.quit();
			}
			break;
		} else {
			cQueue.push(x);
			cQueue.pop();
		}
	}

}

void ThreadPoolAsyncImpl::waitForTask(const PStreamEventDispatcher &sed) noexcept {
			auto t = sed->wait();
			if (t == nullptr) {
				onExitDispatcher(sed);
				return;
			} else {
				dQueue.dispatch(InvokeNetworkDispatcher(*this, sed));
			}
			t();
}

PStreamEventDispatcher ThreadPoolAsyncImpl::getListener() {
	Sync _(lock);
	PStreamEventDispatcher lst,lst2;
	while (cQueue.size() < reqDispatcherCount) {
		lst = AbstractStreamEventDispatcher::create();
		cQueue.push(lst);
		dQueue.dispatch(InvokeNetworkDispatcher(*this, lst));
	}
	while (cQueue.size() > reqDispatcherCount) {
		//pop extra dispatcher
		//they are still in queue, but eventually disappear because they no longer receive a work
		cQueue.pop();
	}
	checkThreadCount();
	lst = cQueue.front();
	cQueue.pop();
	cQueue.push(lst);
	return lst;
}

void ThreadPoolAsyncImpl::runAsync(const AsyncResource& resource, int timeout, const CompletionFn &fn) {

	if (exitFlag) {
		defer >> std::bind(fn, asyncCancel);
		return;
	}

	unsigned int tries = 0;
	auto retry = [&] {
		Sync _(lock);
		++tries;
		if (tries >= cQueue.size())  {
			reqDispatcherCount++;
			if (reqDispatcherCount>reqThreadCount) reqThreadCount++;
		}
	};
	do {
		try {
			auto lst = getListener();
			if (lst->getPendingCount()>=taskLimit) {
				retry();
			} else {
				lst->runAsync(resource,timeout, fn);
				return;
			}
		} catch (OutOfSpaceException &) {
			retry();
		}
	}while(true);
}

void ThreadPoolAsyncImpl::setCountOfDispatchers(unsigned int count) {
	Sync _(lock);
	reqDispatcherCount = count;
}

void ThreadPoolAsyncImpl::setCountOfThreads(unsigned int count) {
	Sync _(lock);
	reqThreadCount = count;
}

void ThreadPoolAsyncImpl::stop() {

	{
		Sync _(lock);
		exitFlag = true;
		auto sz = cQueue.size();
		while (sz) {
			PStreamEventDispatcher d = cQueue.front();
			cQueue.pop();
			cQueue.push(d);
			d->stop();
			--sz;
		}
	}
	threadCount.wait();

}

void ThreadPoolAsyncImpl::setTasksPerDispLimit(unsigned int count) {
	Sync _(lock);
	taskLimit = count;

}

void ThreadPoolAsyncImpl::worker() noexcept {

	using namespace ondra_shared;

	DeferContext defer(ondra_shared::defer_root);

	for(;;) {

		if (!dQueue.pump()) {
			dQueue.quit();
			defer_yield();
			threadCount.dec();
			return;
		}

		defer_yield();
/*		PStreamEventDispatcher lst = tQueue.pop();
		if (lst == nullptr) {
			tQueue.push(nullptr);
			threadCount.dec();
			return;
		}

		auto t = lst->wait();
		tQueue.push(lst);
		if (t == nullptr) {
			threadCount.dec();
			return;
		}

		t();
*/
		if (static_cast<unsigned int>(threadCount.getCounter()) > reqThreadCount) {
			Sync _(lock);
			if (static_cast<unsigned int>(threadCount.getCounter()) > reqThreadCount) {
				threadCount.dec();
				return;
			}
		}
	}
}

AsyncProvider ThreadPoolAsync::create(unsigned int numThreads, unsigned int numListeners, unsigned int tasksPerDispLimit) {
	ThreadPoolAsync provider (new ThreadPoolAsyncImpl);
	provider.setCountOfDispatchers(numListeners);
	provider.setCountOfThreads(numThreads);
	provider.setTasksPerDispLimit(tasksPerDispLimit);
	return provider;
}

void ThreadPoolAsync::setCountOfDispatchers(unsigned int count) {
	ptr->setCountOfDispatchers(count);
}

void ThreadPoolAsync::setCountOfThreads(unsigned int count) {
	ptr->setCountOfThreads(count);
}

void ThreadPoolAsync::stop() {
	ptr->stop();
}

ThreadPoolAsync::~ThreadPoolAsync() {
	if (!ptr->isShared()) {
		ptr->stop();
	}
}

void ThreadPoolAsyncImpl::runAsync(const CustomFn& completion) {

	if (exitFlag) {
		defer >> completion;
		return;
	}

	if (reqThreadCount > reqDispatcherCount) {
		checkThreadCount();
		dQueue.dispatch(completion);
	} else {
		auto lst = getListener();
		lst->runAsync(completion);
	}
}

void ThreadPoolAsync::setTasksPerDispLimit(unsigned int count) {
	ptr->setTasksPerDispLimit(count);
}

ThreadPoolAsync::operator AsyncProvider() const {
	return AsyncProvider(ptr);
}

}

<|endoftext|>