<commit_before>#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <errno.h>
#include <time.h>
#include <limits.h>
#include <ctype.h>
#include <mongoose.h>
#include <json.hpp>

#if defined(_WIN32)
#pragma comment(lib, "ws2_32.lib")
//#pragma comment(lib, "libtesseract302.lib")
//#pragma comment(lib, "liblept168.lib")
#endif

/** SHA256 */
#define SHA256_BLOCK_SIZE 32

typedef struct
{
    unsigned char data[64];
    unsigned int datalen;
    unsigned long long bitlen;
    unsigned int state[8];
} SHA256_CTX;

/*********************** FUNCTION DECLARATIONS **********************/
extern "C" {
    void sha256_init(SHA256_CTX *ctx);
    void sha256_update(SHA256_CTX *ctx, const unsigned char data[], size_t len);
    void sha256_final(SHA256_CTX *ctx, unsigned char hash[]);
}

static void
to_hexstring(unsigned char hash[SHA256_BLOCK_SIZE], char hex[SHA256_BLOCK_SIZE * 2 + 1])
{
    static const char hex_chars[16] = { '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F' };
    for (auto i = 0; i < SHA256_BLOCK_SIZE; ++i) {
        const char b = hash[i];
        hex[i * 2] = hex_chars[(b & 0xf0) >> 4];
        hex[i * 2 + 1] = hex_chars[(b & 0x0f) >> 0];
    }
}
/***********/
static const char *s_http_port = "9090";

struct FileWriterData
{
    FILE      *fd;
    size_t     bytes_left;
    SHA256_CTX ctx;
    char       hash[SHA256_BLOCK_SIZE * 2 + 1];
    char       mime[256];
};

struct FileObject
{
    std::string hash;
    std::string mime;
    std::string content;
    std::string created_at;

    std::string json;
};

char*
time_to_string(const struct tm *timeptr)
{
    static char result[40];
    sprintf(result, "%.2d:%.2d:%.2d %.2d.%.2d.%.4d",
        timeptr->tm_hour,
        timeptr->tm_min,
        timeptr->tm_sec,
        timeptr->tm_mday,
        timeptr->tm_mon + 1,
        1900 + timeptr->tm_year);
    return result;
}

void
recognizer_push(const char * const hash);
void
recognizer_start();

static void
load_file_object(const char * const hash, struct FileObject *fo)
{
    *fo = {};
    std::string filename(hash);
    FILE *fd;
    filename += ".json";
    fd = fopen(filename.c_str(), "r");
    if (fd != nullptr) {
        char        buffer[256];
        std::string content;
        int         ret;

        while ((ret = fread(buffer, 1, sizeof(buffer), fd)) > 0) {
            content += std::string(buffer, ret);
        }
        fclose(fd);

        auto jsonObject = nlohmann::json::parse(content);
        fo->hash = jsonObject["filename"].get<std::string>();
        fo->mime = jsonObject["content_type"].get<std::string>();
        if (!jsonObject["content"].is_null()) {
            fo->content = jsonObject["content"].get<std::string>();
        }
        fo->created_at = jsonObject["created_at"].get<std::string>();
        fo->json = content;
    }
}

static std::string
to_json(const char * const hash, const char * const mime, const char * const content = nullptr, bool store = true)
{
    time_t     rawtime;
    struct tm *timeinfo;

    time(&rawtime);
    timeinfo = localtime(&rawtime);
    nlohmann::json jsonObject = {
        { "filename", hash },
        { "content_type", mime },
        { "created_at", time_to_string(timeinfo) }
    };
    if (content != nullptr)
        jsonObject["content"] = content;
    auto simplify = jsonObject.dump(2);
    if (store) {
        std::string filename(hash);
        FILE *fd;
        filename += ".json";
        fd = fopen(filename.c_str(), "w");
        if (fd != nullptr) {
            fwrite(simplify.c_str(), 1, simplify.size(), fd);
            fclose(fd);
        }
    }
    return simplify;
}

static void
save_object(FileObject &object)
{
    nlohmann::json jsonObject = {
        { "filename", object.hash },
        { "content_type", object.mime },
        { "created_at", object.created_at },
        { "content", object.content },
    };
    std::string filename(object.hash);
    FILE *fd;
    filename += ".json";
    fd = fopen(filename.c_str(), "w");
    auto simplify = jsonObject.dump(2);
    if (fd != nullptr) {
        fwrite(simplify.c_str(), 1, simplify.size(), fd);
        fclose(fd);
    }
}

void
save_content(const char * const hash, const std::string &content)
{
    FileObject object;

    load_file_object(hash, &object);
    object.content = content;
    save_object(object);
}

static void handle_request(struct mg_connection *nc)
{

    mg_printf(nc, "%s",
        "HTTP/1.1 200 OK\r\n"
        "Content-Type: text/html\r\n"
        "Connection: close\r\n"
        "\r\n"
        "Upload example."
        "<form method=\"POST\" action=\"/upload\" "
        "  enctype=\"multipart/form-data\">"
        "<input type=\"file\" name=\"file\" />
"
        "<input type=\"submit\" value=\"Upload\" />"
        "</form>");
    nc->flags |= MG_F_SEND_AND_CLOSE;
}

static void handle_recv(struct mg_connection *nc)
{
    auto data = static_cast<struct FileWriterData *>(nc->user_data);

    if (data == nullptr) {
        // This is a new connection, try to parse HTTP request.
        struct http_message hm;
        int req_len = mg_parse_http(nc->recv_mbuf.buf, nc->recv_mbuf.len, &hm,
            1 /* is_req */);

        if (req_len < 0 ||
            (req_len == 0 && nc->recv_mbuf.len >= MG_MAX_HTTP_REQUEST_SIZE)) {
            nc->flags |= MG_F_CLOSE_IMMEDIATELY;
        } else if (req_len == 0) {
            // Request is not complete yet, do nothing.
        } else if (mg_vcasecmp(&hm.method, "POST") == 0 &&
            mg_vcmp(&hm.uri, "/new") == 0) {
            // This is the request that we don't want to buffer in memory.

            if (hm.body.len == static_cast<size_t>(~0) || hm.body.len == 0) {
                mg_printf(nc, "%s",
                    "HTTP/1.1 411 Content-Length required\r\n"
                    "Content-Length: 0\r\n\r\n");
                nc->flags |= MG_F_SEND_AND_CLOSE;
                return;
            }

            // Reset proto_handler so Mongoose's http_handler() won't get confused
            // once we start removing data from the buffer.
            nc->proto_handler = nullptr;

            // Headers will be inaccessible later, so put everything we need into
            // user_data.
            data = static_cast<struct FileWriterData *>(calloc(1, sizeof(struct FileWriterData)));
            {
                for (auto i = 0; i < MG_MAX_HTTP_HEADERS && hm.header_names[i].p != nullptr; ++i) {
                    auto n = hm.header_names[i];
                    if (!strnicmp(n.p, "Content-Type", n.len)) {
                        n = hm.header_values[i];
                        auto it = n.p;
                        auto len = n.len - (it - n.p);
                        while (it && *it == ' ' && (it - n.p) < n.len)
                            ++it;
                        auto start = it;
                        it = n.p + n.len - 1;
                        while (it && (*it == '\r' || *it == '\n' || *it == ' ') && it > n.p)
                            --it, --len;
                        strncpy(data->mime, start, len);
                    }
                }
            }
            data->bytes_left = hm.body.len;
            data->fd = tmpfile();
            sha256_init(&data->ctx);
            if (data->fd == nullptr) {
                mg_printf(nc, "%s",
                    "HTTP/1.1 500 Failed to open a file\r\n"
                    "Content-Length: 0\r\n\r\n");
                nc->flags |= MG_F_SEND_AND_CLOSE;
                return;
            }
            nc->user_data = static_cast<void *>(data);

            // Remove the headers from the buffer.
            mbuf_remove(&nc->recv_mbuf, hm.body.p - nc->recv_mbuf.buf);

            // Invoke itself again to write the piece of the body that is already in
            // the buffer.
            handle_recv(nc);
        }
    } else {  // data != NULL
        size_t to_write = data->bytes_left, written = 0;
        if (nc->recv_mbuf.len < to_write)
            to_write = nc->recv_mbuf.len;
        written = fwrite(nc->recv_mbuf.buf, 1, to_write, data->fd);
        sha256_update(&data->ctx, reinterpret_cast<unsigned char * const>(nc->recv_mbuf.buf), written);
        mbuf_remove(&nc->recv_mbuf, written);
        data->bytes_left -= written;
        if (data->bytes_left <= 0) {
            FileObject object;
            std::string json;
            long        size;
            unsigned char digest[SHA256_BLOCK_SIZE];

            // Request is complete, do something meaningful here.
            sha256_final(&data->ctx, digest);
            to_hexstring(digest, data->hash);
            load_file_object(data->hash, &object);
            {
                size = ftell(data->fd);
                fseek(data->fd, 0, SEEK_SET);
                auto fd = fopen(data->hash, "wb");
                if (fd) {
                    char buffer[256];
                    int  readed;
                    while ((readed = fread(buffer, 1, sizeof(buffer), data->fd)) > 0) {
                        fwrite(buffer, 1, readed, fd);
                    }
                    fclose(fd);
                }
            }
            if (!object.hash.empty()) {
                fprintf(stdout, "%s already stored at %s\n", data->hash, object.created_at.c_str());
                json = object.json;
            } else {
                json = to_json(data->hash, data->mime);
                recognizer_push(data->hash);
            }
            mg_printf(nc,
                "HTTP/1.1 200 OK\r\n"
                "Content-Type: application/json; charset=utf-8\r\n"
                "Connection: close\r\n\r\n"
                "%s",
                json.c_str());
            nc->flags |= MG_F_SEND_AND_CLOSE;
            // handle_close will free the resources.
        }
    }
}

// Make sure we free all allocated resources
static void handle_close(struct mg_connection *nc)
{
    auto data = static_cast<struct FileWriterData *>(nc->user_data);

    if (data != nullptr) {

        fclose(data->fd);
        free(data);
        nc->user_data = nullptr;
    }
}

static void ev_handler(struct mg_connection *nc, int ev, void *ev_data)
{
    (void)ev_data;
    switch (ev) {
    case MG_EV_HTTP_REQUEST:
        // Invoked when the full HTTP request is in the buffer (including body).
        handle_request(nc);
        break;
    case MG_EV_RECV:
        // Invoked every time new data arrives.
        handle_recv(nc);
        break;
    case MG_EV_CLOSE:
        handle_close(nc);
        break;
    }
}

int main(void)
{
    struct mg_mgr mgr;
    struct mg_connection *nc;

    mg_mgr_init(&mgr, nullptr);
    nc = mg_bind(&mgr, s_http_port, ev_handler);

    // Set up HTTP server parameters
    mg_set_protocol_http_websocket(nc);
    mg_enable_multithreading(nc);

    printf("Starting web server on port %s\n", s_http_port);
    recognizer_start();
    for (;;) {
        mg_mgr_poll(&mgr, 1000);
    }
    mg_mgr_free(&mgr);

    return EXIT_SUCCESS;
}
<commit_msg>FIX: strnicmp on linux<commit_after>#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <errno.h>
#include <time.h>
#include <limits.h>
#include <ctype.h>
#include <mongoose.h>
#include <json.hpp>

#if defined(_WIN32)
#pragma comment(lib, "ws2_32.lib")
//#pragma comment(lib, "libtesseract302.lib")
//#pragma comment(lib, "liblept168.lib")
#else
#define strnicmp strncasecmp
#endif

/** SHA256 */
#define SHA256_BLOCK_SIZE 32

typedef struct
{
    unsigned char data[64];
    unsigned int datalen;
    unsigned long long bitlen;
    unsigned int state[8];
} SHA256_CTX;

/*********************** FUNCTION DECLARATIONS **********************/
extern "C" {
    void sha256_init(SHA256_CTX *ctx);
    void sha256_update(SHA256_CTX *ctx, const unsigned char data[], size_t len);
    void sha256_final(SHA256_CTX *ctx, unsigned char hash[]);
}

static void
to_hexstring(unsigned char hash[SHA256_BLOCK_SIZE], char hex[SHA256_BLOCK_SIZE * 2 + 1])
{
    static const char hex_chars[16] = { '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F' };
    for (auto i = 0; i < SHA256_BLOCK_SIZE; ++i) {
        const char b = hash[i];
        hex[i * 2] = hex_chars[(b & 0xf0) >> 4];
        hex[i * 2 + 1] = hex_chars[(b & 0x0f) >> 0];
    }
}
/***********/
static const char *s_http_port = "9090";

struct FileWriterData
{
    FILE      *fd;
    size_t     bytes_left;
    SHA256_CTX ctx;
    char       hash[SHA256_BLOCK_SIZE * 2 + 1];
    char       mime[256];
};

struct FileObject
{
    std::string hash;
    std::string mime;
    std::string content;
    std::string created_at;

    std::string json;
};

char*
time_to_string(const struct tm *timeptr)
{
    static char result[40];
    sprintf(result, "%.2d:%.2d:%.2d %.2d.%.2d.%.4d",
        timeptr->tm_hour,
        timeptr->tm_min,
        timeptr->tm_sec,
        timeptr->tm_mday,
        timeptr->tm_mon + 1,
        1900 + timeptr->tm_year);
    return result;
}

void
recognizer_push(const char * const hash);
void
recognizer_start();

static void
load_file_object(const char * const hash, struct FileObject *fo)
{
    *fo = {};
    std::string filename(hash);
    FILE *fd;
    filename += ".json";
    fd = fopen(filename.c_str(), "r");
    if (fd != nullptr) {
        char        buffer[256];
        std::string content;
        int         ret;

        while ((ret = fread(buffer, 1, sizeof(buffer), fd)) > 0) {
            content += std::string(buffer, ret);
        }
        fclose(fd);

        auto jsonObject = nlohmann::json::parse(content);
        fo->hash = jsonObject["filename"].get<std::string>();
        fo->mime = jsonObject["content_type"].get<std::string>();
        if (!jsonObject["content"].is_null()) {
            fo->content = jsonObject["content"].get<std::string>();
        }
        fo->created_at = jsonObject["created_at"].get<std::string>();
        fo->json = content;
    }
}

static std::string
to_json(const char * const hash, const char * const mime, const char * const content = nullptr, bool store = true)
{
    time_t     rawtime;
    struct tm *timeinfo;

    time(&rawtime);
    timeinfo = localtime(&rawtime);
    nlohmann::json jsonObject = {
        { "filename", hash },
        { "content_type", mime },
        { "created_at", time_to_string(timeinfo) }
    };
    if (content != nullptr)
        jsonObject["content"] = content;
    auto simplify = jsonObject.dump(2);
    if (store) {
        std::string filename(hash);
        FILE *fd;
        filename += ".json";
        fd = fopen(filename.c_str(), "w");
        if (fd != nullptr) {
            fwrite(simplify.c_str(), 1, simplify.size(), fd);
            fclose(fd);
        }
    }
    return simplify;
}

static void
save_object(FileObject &object)
{
    nlohmann::json jsonObject = {
        { "filename", object.hash },
        { "content_type", object.mime },
        { "created_at", object.created_at },
        { "content", object.content },
    };
    std::string filename(object.hash);
    FILE *fd;
    filename += ".json";
    fd = fopen(filename.c_str(), "w");
    auto simplify = jsonObject.dump(2);
    if (fd != nullptr) {
        fwrite(simplify.c_str(), 1, simplify.size(), fd);
        fclose(fd);
    }
}

void
save_content(const char * const hash, const std::string &content)
{
    FileObject object;

    load_file_object(hash, &object);
    object.content = content;
    save_object(object);
}

static void handle_request(struct mg_connection *nc)
{

    mg_printf(nc, "%s",
        "HTTP/1.1 200 OK\r\n"
        "Content-Type: text/html\r\n"
        "Connection: close\r\n"
        "\r\n"
        "Upload example."
        "<form method=\"POST\" action=\"/upload\" "
        "  enctype=\"multipart/form-data\">"
        "<input type=\"file\" name=\"file\" />
"
        "<input type=\"submit\" value=\"Upload\" />"
        "</form>");
    nc->flags |= MG_F_SEND_AND_CLOSE;
}

static void handle_recv(struct mg_connection *nc)
{
    auto data = static_cast<struct FileWriterData *>(nc->user_data);

    if (data == nullptr) {
        // This is a new connection, try to parse HTTP request.
        struct http_message hm;
        int req_len = mg_parse_http(nc->recv_mbuf.buf, nc->recv_mbuf.len, &hm,
            1 /* is_req */);

        if (req_len < 0 ||
            (req_len == 0 && nc->recv_mbuf.len >= MG_MAX_HTTP_REQUEST_SIZE)) {
            nc->flags |= MG_F_CLOSE_IMMEDIATELY;
        } else if (req_len == 0) {
            // Request is not complete yet, do nothing.
        } else if (mg_vcasecmp(&hm.method, "POST") == 0 &&
            mg_vcmp(&hm.uri, "/new") == 0) {
            // This is the request that we don't want to buffer in memory.

            if (hm.body.len == static_cast<size_t>(~0) || hm.body.len == 0) {
                mg_printf(nc, "%s",
                    "HTTP/1.1 411 Content-Length required\r\n"
                    "Content-Length: 0\r\n\r\n");
                nc->flags |= MG_F_SEND_AND_CLOSE;
                return;
            }

            // Reset proto_handler so Mongoose's http_handler() won't get confused
            // once we start removing data from the buffer.
            nc->proto_handler = nullptr;

            // Headers will be inaccessible later, so put everything we need into
            // user_data.
            data = static_cast<struct FileWriterData *>(calloc(1, sizeof(struct FileWriterData)));
            {
                for (auto i = 0; i < MG_MAX_HTTP_HEADERS && hm.header_names[i].p != nullptr; ++i) {
                    auto n = hm.header_names[i];
                    if (!strnicmp(n.p, "Content-Type", n.len)) {
                        n = hm.header_values[i];
                        auto it = n.p;
                        auto len = n.len - (it - n.p);
                        while (it && *it == ' ' && (it - n.p) < n.len)
                            ++it;
                        auto start = it;
                        it = n.p + n.len - 1;
                        while (it && (*it == '\r' || *it == '\n' || *it == ' ') && it > n.p)
                            --it, --len;
                        strncpy(data->mime, start, len);
                    }
                }
            }
            data->bytes_left = hm.body.len;
            data->fd = tmpfile();
            sha256_init(&data->ctx);
            if (data->fd == nullptr) {
                mg_printf(nc, "%s",
                    "HTTP/1.1 500 Failed to open a file\r\n"
                    "Content-Length: 0\r\n\r\n");
                nc->flags |= MG_F_SEND_AND_CLOSE;
                return;
            }
            nc->user_data = static_cast<void *>(data);

            // Remove the headers from the buffer.
            mbuf_remove(&nc->recv_mbuf, hm.body.p - nc->recv_mbuf.buf);

            // Invoke itself again to write the piece of the body that is already in
            // the buffer.
            handle_recv(nc);
        }
    } else {  // data != NULL
        size_t to_write = data->bytes_left, written = 0;
        if (nc->recv_mbuf.len < to_write)
            to_write = nc->recv_mbuf.len;
        written = fwrite(nc->recv_mbuf.buf, 1, to_write, data->fd);
        sha256_update(&data->ctx, reinterpret_cast<unsigned char * const>(nc->recv_mbuf.buf), written);
        mbuf_remove(&nc->recv_mbuf, written);
        data->bytes_left -= written;
        if (data->bytes_left <= 0) {
            FileObject object;
            std::string json;
            long        size;
            unsigned char digest[SHA256_BLOCK_SIZE];

            // Request is complete, do something meaningful here.
            sha256_final(&data->ctx, digest);
            to_hexstring(digest, data->hash);
            load_file_object(data->hash, &object);
            {
                size = ftell(data->fd);
                fseek(data->fd, 0, SEEK_SET);
                auto fd = fopen(data->hash, "wb");
                if (fd) {
                    char buffer[256];
                    int  readed;
                    while ((readed = fread(buffer, 1, sizeof(buffer), data->fd)) > 0) {
                        fwrite(buffer, 1, readed, fd);
                    }
                    fclose(fd);
                }
            }
            if (!object.hash.empty()) {
                fprintf(stdout, "%s already stored at %s\n", data->hash, object.created_at.c_str());
                json = object.json;
            } else {
                json = to_json(data->hash, data->mime);
                recognizer_push(data->hash);
            }
            mg_printf(nc,
                "HTTP/1.1 200 OK\r\n"
                "Content-Type: application/json; charset=utf-8\r\n"
                "Connection: close\r\n\r\n"
                "%s",
                json.c_str());
            nc->flags |= MG_F_SEND_AND_CLOSE;
            // handle_close will free the resources.
        }
    }
}

// Make sure we free all allocated resources
static void handle_close(struct mg_connection *nc)
{
    auto data = static_cast<struct FileWriterData *>(nc->user_data);

    if (data != nullptr) {

        fclose(data->fd);
        free(data);
        nc->user_data = nullptr;
    }
}

static void ev_handler(struct mg_connection *nc, int ev, void *ev_data)
{
    (void)ev_data;
    switch (ev) {
    case MG_EV_HTTP_REQUEST:
        // Invoked when the full HTTP request is in the buffer (including body).
        handle_request(nc);
        break;
    case MG_EV_RECV:
        // Invoked every time new data arrives.
        handle_recv(nc);
        break;
    case MG_EV_CLOSE:
        handle_close(nc);
        break;
    }
}

int main(void)
{
    struct mg_mgr mgr;
    struct mg_connection *nc;

    mg_mgr_init(&mgr, nullptr);
    nc = mg_bind(&mgr, s_http_port, ev_handler);

    // Set up HTTP server parameters
    mg_set_protocol_http_websocket(nc);
    mg_enable_multithreading(nc);

    printf("Starting web server on port %s\n", s_http_port);
    recognizer_start();
    for (;;) {
        mg_mgr_poll(&mgr, 1000);
    }
    mg_mgr_free(&mgr);

    return EXIT_SUCCESS;
}
<|endoftext|>


<commit_before>#include <iostream>

#include "player.h"
#include "Hand.h"

const char* Player::VERSION = "Non-Default C++ folding player";

int minimum_raise;
int in_action;
int our_bet;
int current_buy_in;
Hand2 hole_cards;

int toAction(HoleCardRank rank) {
    std::cerr << "XXXXXXXXX action is " << int(rank) << std::endl;

    switch (rank) {
        case HoleCardRank::ALLIN:
            return 5000;
        case HoleCardRank::RAISABLE:
            return current_buy_in - our_bet + 3*minimum_raise;
        default:
            std::cerr << "Unknown action" << std::endl;
        case HoleCardRank::FOLDABLE:
            return 0;
        case HoleCardRank::CALLABLE:
            return current_buy_in - our_bet;
    }
}

int preFlop(json::Value game_state) {
    return toAction(rankHoleCard(hole_cards));
}

int flop(json::Value game_state) {
    if (!game_state.HasKey("eval")) {
        std::cerr << "NO eval in flop()" << std::endl;
        return preFlop(game_state);
    }

    int rank = game_state["eval"]["rank"].ToInt();

    switch (rank) {
        default:
            std::cerr << "Unknown rank" << std::endl;
        case 0: return toAction(HoleCardRank::FOLDABLE);
        case 1: return toAction(HoleCardRank::CALLABLE);
        case 2: return toAction(HoleCardRank::CALLABLE);
        case 3: return toAction(HoleCardRank::RAISABLE);
        case 4: return toAction(HoleCardRank::RAISABLE);
        case 5: return toAction(HoleCardRank::ALLIN);
        case 6: return toAction(HoleCardRank::ALLIN);
        case 7: return toAction(HoleCardRank::ALLIN);
        case 8: return toAction(HoleCardRank::ALLIN);
    }
}

int turn(json::Value game_state) {
    if (!game_state.HasKey("eval")) {
        std::cerr << "NO eval in turn()" << std::endl;
        return preFlop(game_state);
    }
    return flop(game_state);
}

int river(json::Value game_state) {
    if (!game_state.HasKey("eval")) {
        std::cerr << "NO eval in river()" << std::endl;
        return preFlop(game_state);
    }
    return flop(game_state);
}

int Player::betRequest(json::Value game_state) {
    minimum_raise = game_state["minimum_raise"].ToInt();
    in_action = game_state["in_action"].ToInt();
    our_bet = game_state["players"][in_action]["bet"];
    current_buy_in = game_state["current_buy_in"];

    hole_cards = parseHand(game_state["players"][in_action]["hole_cards"]);

    std::cerr << "XXXXXXXXXX Our hand is = " << hole_cards << std::endl;

    switch (game_state["community_cards"].ToArray().size()) {
        default:
        case 0:
            return preFlop(game_state);
        case 3:
            return flop(game_state);
        case 4:
            return turn(game_state);
        case 5:
            return river(game_state);
    }

}

void Player::showdown(json::Value game_state)
{
}
<commit_msg>Pwnies<commit_after>#include <iostream>

#include "player.h"
#include "Hand.h"

const char* Player::VERSION = "Pwnies";

int minimum_raise;
int in_action;
int our_bet;
int current_buy_in;
Hand2 hole_cards;

int toAction(HoleCardRank rank) {
    std::cerr << "XXXXXXXXX action is " << int(rank) << std::endl;

    switch (rank) {
        case HoleCardRank::ALLIN:
            return 5000;
        case HoleCardRank::RAISABLE:
            return current_buy_in - our_bet + 3*minimum_raise;
        default:
            std::cerr << "Unknown action" << std::endl;
        case HoleCardRank::FOLDABLE:
            return 0;
        case HoleCardRank::CALLABLE:
            return current_buy_in - our_bet;
    }
}

int preFlop(json::Value game_state) {
    return toAction(rankHoleCard(hole_cards));
}

int flop(json::Value game_state) {
    if (!game_state.HasKey("eval")) {
        std::cerr << "NO eval in flop()" << std::endl;
        return preFlop(game_state);
    }

    int rank = game_state["eval"]["rank"].ToInt();

    switch (rank) {
        default:
            std::cerr << "Unknown rank" << std::endl;
        case 0: return toAction(HoleCardRank::FOLDABLE);
        case 1: return toAction(HoleCardRank::CALLABLE);
        case 2: return toAction(HoleCardRank::CALLABLE);
        case 3: return toAction(HoleCardRank::RAISABLE);
        case 4: return toAction(HoleCardRank::RAISABLE);
        case 5: return toAction(HoleCardRank::ALLIN);
        case 6: return toAction(HoleCardRank::ALLIN);
        case 7: return toAction(HoleCardRank::ALLIN);
        case 8: return toAction(HoleCardRank::ALLIN);
    }
}

int turn(json::Value game_state) {
    if (!game_state.HasKey("eval")) {
        std::cerr << "NO eval in turn()" << std::endl;
        return preFlop(game_state);
    }
    return flop(game_state);
}

int river(json::Value game_state) {
    if (!game_state.HasKey("eval")) {
        std::cerr << "NO eval in river()" << std::endl;
        return preFlop(game_state);
    }
    return flop(game_state);
}

int Player::betRequest(json::Value game_state) {
    minimum_raise = game_state["minimum_raise"].ToInt();
    in_action = game_state["in_action"].ToInt();
    our_bet = game_state["players"][in_action]["bet"];
    current_buy_in = game_state["current_buy_in"];

    hole_cards = parseHand(game_state["players"][in_action]["hole_cards"]);

    std::cerr << "XXXXXXXXXX Our hand is = " << hole_cards << std::endl;

    switch (game_state["community_cards"].ToArray().size()) {
        default:
        case 0:
            return preFlop(game_state);
        case 3:
            return flop(game_state);
        case 4:
            return turn(game_state);
        case 5:
            return river(game_state);
    }

}

void Player::showdown(json::Value game_state)
{
}
<|endoftext|>


<commit_before>#include <cmath>
#include <cstdlib>
#include <ctime>
#include "Experience.h"
#include "StateSpace.h"
#include "Action.h"

//utility calculation function
double q();

//reward function based off raw position and velocity
double state_reward(const double theta, const double theta_dot);

int main()
{
	std::srand(std::time(NULL));

	//create the state space
	StateSpace space(100,50);

	//array to hold the current state of the robot
	double current_state[2]={0,0};
	double last_state[2]={0,0};

	//create pointers to the possible actions as well as pointers to hold the chosen action and the previous action
	Action* actions[2],chosen_action;
	actions[0]=new action(FORWARDS);
	actions[1]=new action(BACKWARDS);

	//evaluate state
	current_state[0]=getPosition();//!
	current_state[1]=getVelocity();//!

	while(true)
	{
		//evaluate q of last action
		//current_state>>old_state

		chosen_action=selectAction(/*...*/);

		//>>current_state
	}

	delete actions[0];
	delete actions[1];

	return 1;
}

double q()
{
	return //StateSpace[originalstate[0],originalstate[1]].GetExperience().GetUtility(action) + learningrate(state_reward(newstate[0], newstate[1]) + gamma(StateSpace[newstate[0],newstate[1]].GetExperience().GetHighestUtility()) - StateSpace[originalstate[0],originalstate[1]].GetExperience().GetUtility(action); // Utility calculation pseudocode
}

double state_reward(const double theta, const double theta_dot)
{
	return 180+theta_dot*theta/std::abs(theta_dot);
}

//function is fed with a priority queue of action-values
//generates Boltzmann distribution of these action-values
//and selects an action based on probabilities
Action * selectAction(PriorityQueue<Action *,double> a_queue, double temp)
{
    typedef PriorityQueue<Action *,double> PQ;
    typedef std::vector< std::pair<Action *, double> > Vec_Pair;
    typedef std::pair<Action *, double> Pair;

    double sum= 0;
    int i = 0;

    int size = a_queue.getSize();
    Vec_Pair action_vec(size);
    Pair pear;

    //Calculate partition function by iterating over action-values
	for(PQ::constIter iter = a_queue.begin(); iter < a_queue.end(); iter++)
	{
		sum += std::exp((iter->second)/temp);
	}
	//Calculate boltzmann factors for action-values
	for(Vec_Pair::iterator it = action_vec.begin(); it < action_vec.end(); it++)
	{
	    it->first = a_queue[i].first;
	    it->second = std::exp(a_queue[i].second / temp) / sum;
	    i++;
	}
	//calculate cumulative probability distribution
    for(std::vector< std::pair<int, double> >::iterator it = action_vec.begin()+1; it < action_vec.end(); it++)
	{
	    it->second += (it-1)->second;
	}
    //seed random number generator
	srand(time(NULL));
	//generate RN between 0 and 1
	double rand_num = (double)rand()/ (RAND_MAX);

	//select action based on probability
    for(Vec_Pair::iterator it = action_vec.begin(); it < action_vec.end(); it++)
    {
        //if RN falls within cumulative probability bin return the corresponding action
        if(rand_num < it->second){return it->first};
    }
    return -1;
}
<commit_msg>added Q update function<commit_after>#include <cmath>
#include <cstdlib>
#include <ctime>
#include "Experience.h"
#include "StateSpace.h"
#include "Action.h"

//utility calculation function
double q();

//reward function based off raw position and velocity
double state_reward(const double theta, const double theta_dot);

int main()
{
	std::srand(std::time(NULL));

	//create the state space
	StateSpace space(100,50);

	//array to hold the current state of the robot
	double current_state[2]={0,0};
	double last_state[2]={0,0};

	//create pointers to the possible actions as well as pointers to hold the chosen action and the previous action
	Action* actions[2],chosen_action;
	actions[0]=new action(FORWARDS);
	actions[1]=new action(BACKWARDS);

	//evaluate state
	current_state[0]=getPosition();//!
	current_state[1]=getVelocity();//!

	while(true)
	{
		//evaluate q of last action
		//current_state>>old_state

		chosen_action=selectAction(/*...*/);

		//>>current_state
	}

	delete actions[0];
	delete actions[1];

	return 1;
}

double q()
{
	return //StateSpace[originalstate[0],originalstate[1]].GetExperience().GetUtility(action) + learningrate(state_reward(newstate[0], newstate[1]) + gamma(StateSpace[newstate[0],newstate[1]].GetExperience().GetHighestUtility()) - StateSpace[originalstate[0],originalstate[1]].GetExperience().GetUtility(action); // Utility calculation pseudocode
}

double state_reward(const double theta, const double theta_dot)
{
	return 180+theta_dot*theta/std::abs(theta_dot);
}
//updates Q value for a state-action
void updateQ(StateSpace * space, Action * action, State * current_state,
                 State * old_state, double alpha, double gamma)
{
    //oldQ value
    double oldQ = (*space).StateSearch(old_state).search(action).second;
    //reward given to current state
    double R = new_state->getReward();
    //optimal Q value for new state i.e. first element
    double maxQ = (*space).StateSearch(current_state)[0].second;

    //new Q value determined by Q learning algorithm
    double newQ = oldQ + alpha * (R + (gamma * maxQ) - oldQ;

    //updates Q value
    (*space).StateSearch(old_state).search(action).second = newQ;
}

//function is fed with a priority queue of action-values
//generates Boltzmann distribution of these action-values
//and selects an action based on probabilities
Action * selectAction(PriorityQueue<Action *,double> a_queue, double temp)
{
    typedef PriorityQueue<Action *,double> PQ;
    typedef std::vector< std::pair<Action *, double> > Vec_Pair;
    typedef std::pair<Action *, double> Pair;

    double sum= 0;
    int i = 0;

    int size = a_queue.getSize();
    Vec_Pair action_vec(size);
    Pair pear;

    //Calculate partition function by iterating over action-values
	for(PQ::constIter iter = a_queue.begin(); iter < a_queue.end(); iter++)
	{
		sum += std::exp((iter->second)/temp);
	}
	//Calculate boltzmann factors for action-values
	for(Vec_Pair::iterator it = action_vec.begin(); it < action_vec.end(); it++)
	{
	    it->first = a_queue[i].first;
	    it->second = std::exp(a_queue[i].second / temp) / sum;
	    i++;
	}
	//calculate cumulative probability distribution
    for(std::vector< std::pair<int, double> >::iterator it = action_vec.begin()+1; it < action_vec.end(); it++)
	{
	    it->second += (it-1)->second;
	}
    //seed random number generator
	srand(time(NULL));
	//generate RN between 0 and 1
	double rand_num = (double)rand()/ (RAND_MAX);

	//select action based on probability
    for(Vec_Pair::iterator it = action_vec.begin(); it < action_vec.end(); it++)
    {
        //if RN falls within cumulative probability bin return the corresponding action
        if(rand_num < it->second){return it->first};
    }
    return -1;
}
<|endoftext|>


<commit_before>#include "../../wrdTest.hpp"

using namespace wrd;
using namespace std;

struct primitiveObjTest : public wrdTest {};

TEST_F(primitiveObjTest, testCreateWIntInMgd) {
    wInt origin(1);

    tstr<wInt> inst = origin.run(narr({origin}));
    ASSERT_TRUE(inst);
    ASSERT_EQ(origin.get(), inst->get());
    ASSERT_NE(&origin, &(*inst));
}

TEST_F(primitiveObjTest, testCloneWIntInMgd) {
    wInt origin(1);

    tstr<wInt> inst(origin.clone());
    ASSERT_TRUE(inst);
    ASSERT_EQ(origin.get(), inst->get());
    ASSERT_NE(&origin, &(*inst));
}

TEST_F(primitiveObjTest, testDefaultCtor) {
    wInt origin(1);

    tstr<wInt> inst = origin.run(narr());
    ASSERT_TRUE(inst);
    ASSERT_EQ(inst->get(), 0);
}
<commit_msg>wrd: test: verify primitive types uses dummy subs<commit_after>#include "../../wrdTest.hpp"

using namespace wrd;
using namespace std;

struct primitiveObjTest : public wrdTest {};

TEST_F(primitiveObjTest, testCreateWIntInMgd) {
    wInt origin(1);

    tstr<wInt> inst = origin.run(narr({origin}));
    ASSERT_TRUE(inst);
    ASSERT_EQ(origin.get(), inst->get());
    ASSERT_NE(&origin, &(*inst));
}

TEST_F(primitiveObjTest, testCloneWIntInMgd) {
    wInt origin(1);

    tstr<wInt> inst(origin.clone());
    ASSERT_TRUE(inst);
    ASSERT_EQ(origin.get(), inst->get());
    ASSERT_NE(&origin, &(*inst));
}

TEST_F(primitiveObjTest, testDefaultCtor) {
    wInt origin(1);

    tstr<wInt> inst = origin.run(narr());
    ASSERT_TRUE(inst);
    ASSERT_EQ(inst->get(), 0);
}

TEST_F(primitiveObjTest, subsIsDummy) {
    wInt val(2);
    ASSERT_EQ(val.subs().len(), 0);

    bicontainable& subs = val.subs();
    ASSERT_FALSE(nul(subs));
    subs.add("not work", new wInt(3));
    ASSERT_EQ(subs.len(), 0);
    ASSERT_FALSE(subs.has("not work"));
    ASSERT_TRUE(subs.begin().isEnd());
}
<|endoftext|>


<commit_before>#if !defined(__CINT__) || defined(__MAKECINT__)
   #include <TMath.h>
   #include <TROOT.h>
   #include <Riostream.h>
   #include <TCanvas.h>

   #include <TString.h>

   #include <TFile.h>
   #include <TList.h>
   #include <TH1F.h>
   #include <TH1D.h>
   #include <TF2.h>
   #include <TFitResult.h>
   #include <TFitResultPtr.h>
   #include <TH2F.h>
   #include <TH3F.h>
#endif

extern TROOT *gROOT;
//extern TFile *gFile;

void PostProcessCTau() {
  //-----------------------------------------------------------------
  // PostProcess the histograms produced by AliAnalysisTaskCTau*
  // Produce:
  // -- V0 particle spectra (not corrected for the feeddown)
  // -- Estimated (anti)Lambda feed-down spectra (from Xi's only)
  // -- V0 particle c*tau distributions (feeddown corrected)
  //-----------------------------------------------------------------

  TString name("cTau_0090"); //centrality
  //TString name("cTau_0005"); //centrality
  //TString name("cTau_2040"); //centrality
  //TString name("cTau_4060"); //centrality
  //TString name("cTau_6080"); //centrality
  //TString name("cTau_8090"); //centrality

  //+++++ Real data histograms
  TFile::Open("new/real/all/" + name + ".root");
  TList *list=(TList *)gFile->Get(name);

  TH2F *fK0sSi=(TH2F*)list->FindObject("fK0sSi");fK0sSi->Sumw2();
  TH2F *fLambdaSi=(TH2F*)list->FindObject("fLambdaSi");fLambdaSi->Sumw2();
  TH1F *fXiSiP=(TH1F*)list->FindObject("fXiSiP"); fXiSiP->Sumw2();
  TH2F *fLambdaBarSi=
        (TH2F*)list->FindObject("fLambdaBarSi"); fLambdaBarSi->Sumw2();
  TH1F *fXiBarSiP=(TH1F*)list->FindObject("fXiBarSiP"); fXiBarSiP->Sumw2();

  TH1F *fMult=(TH1F*)list->FindObject("fMult");
  //+++++

  const Double_t nEvents=fMult->Integral();

  //+++++ MC histograms
  TFile::Open("new/mc_b_plus/all/" + name + "_mc.root");
  TList *listmc=(TList *)gFile->Get(name+"_mc");

  TH2F *fK0sMC=(TH2F*)listmc->FindObject("fK0sMC"); fK0sMC->Sumw2();
  TH2F *fK0sAs=(TH2F*)listmc->FindObject("fK0sAs"); fK0sAs->Sumw2();

  TH2F *fLambdaMC=(TH2F*)listmc->FindObject("fLambdaMC"); fLambdaMC->Sumw2();
  TH2F *fLambdaAs=(TH2F*)listmc->FindObject("fLambdaAs"); fLambdaAs->Sumw2();

  TH2F *fLambdaBarMC=
        (TH2F*)listmc->FindObject("fLambdaBarMC"); fLambdaBarMC->Sumw2();
  TH2F *fLambdaBarAs=
        (TH2F*)listmc->FindObject("fLambdaBarAs"); fLambdaBarAs->Sumw2();

  TH1F *fXiSiPMC=(TH1F*)listmc->FindObject("fXiSiP");fXiSiPMC->Sumw2();
  TH3F *fLambdaFromXi=(TH3F*)listmc->FindObject("fLambdaFromXi");
  fLambdaFromXi->Sumw2();

  TH1F *fXiBarSiPMC=(TH1F*)listmc->FindObject("fXiBarSiP");fXiBarSiPMC->Sumw2();
  TH3F *fLambdaBarFromXiBar=(TH3F*)listmc->FindObject("fLambdaBarFromXiBar");
  fLambdaBarFromXiBar->Sumw2();
  //+++++

  Double_t myExp2(Double_t *v, Double_t *p);
  void Correct(TH1 *rw, const TH1 *as, const TH1 *mc);
  void Normalise(Double_t, Double_t, Double_t, Double_t, TH1 *h);

  const Int_t    iMax=3;   // Number of V0 particles
  const TString  pnam[]={"K^{0}_{S}", "#Lambda", "#bar{#Lambda}"};
  const Double_t brch[]={0.69, 0.64, 0.64};
  const Double_t mass[]={0.4977, 1.115, 1.115};
  const Double_t ctau[]={2.68, 7.89, 7.89};

  const TH2 *in[]={
    fK0sSi,fK0sAs,fK0sMC,
    fLambdaSi,fLambdaAs,fLambdaMC,
    fLambdaBarSi,fLambdaBarAs,fLambdaBarMC
  };
  TH1 *fd[]={
    0,0,0,
    fLambdaFromXi,fXiSiP,fXiSiPMC,
    fLambdaBarFromXiBar,fXiBarSiP,fXiBarSiPMC
  };
  Double_t wbx=fK0sSi->GetBinWidth(3);
  Int_t nbx=fLambdaFromXi->GetNbinsX();
  Int_t nby=fLambdaFromXi->GetNbinsY();
  Int_t nbz=fLambdaFromXi->GetNbinsZ();

  for (Int_t i=0; i<iMax; i++) {
      TH2 *cr=(TH2*)in[i*3 + 0]->Clone();
      Correct(cr, in[i*3 + 1], in[i*3 + 2]);

      //++++ pT spectrum
      TH1 *pt=cr->ProjectionX("_px",0,-1,"e");
      TString ptName = pnam[i] + " p_{T} spectrum";
      pt->SetTitle(ptName.Data());
      Normalise(brch[i], 2*0.5, wbx, nEvents, pt);
      Double_t eipt=0., ipt=pt->IntegralAndError(1, nbx, eipt);

      new TCanvas;
      pt->Draw();

      if (i>0) {
      //++++ feeddown
	 TH3 *fd3=(TH3*)fd[3*i + 0];
	 TH1 *rl =(TH1*)fd[3*i + 1];
	 TH1 *mc =(TH1*)fd[3*i + 2];
         rl->Divide(mc);

         for (Int_t k=1; k<=nbx; k++) {
             for (Int_t l=1; l<=nby; l++) {
                 for (Int_t m=1; m<=nbz; m++) {
                     Float_t c=fd3->GetBinContent(k,l,m);
                     c *= rl->GetBinContent(m);
                     fd3->SetBinContent(k,l,m,c);
		 }
	     }
	 }

         TH2 *fd2=(TH2*)fd3->Project3D("yxe");
         Correct(fd2, in[i*3 + 1], in[i*3 + 2]);

         TH1 *fd1=fd2->ProjectionX("_px",0,-1,"e");
         Normalise(brch[i], 2*0.5, wbx, nEvents, fd1);
         Double_t eifd=0., ifd=fd1->IntegralAndError(1, nbx, eifd);

         Double_t f=ifd/ipt;
         Double_t ef=1/ipt*TMath::Sqrt(eifd*eifd + f*f*eipt*eipt);
         cout<<endl<<"Global FD correction: "<<f<<"+/-"<<ef<<endl;

         //new TCanvas();
         fd1->Draw("same");

         cr->Add(fd2,-1);
      }
      //continue;

      //++++ c*tau
      TF2 *f2=new TF2("myexpo2",myExp2,0.,10.,0.,100.,1+1+1+nbx);
      f2->SetParameter(0, ctau[i]);
      f2->FixParameter(1, ctau[i]);
      f2->FixParameter(2, mass[i]);
      for (Int_t p=1+1+1; p<=nbx+1+1; p++)
          f2->SetParameter(p,fK0sSi->GetBinContent(p,1));

      new TCanvas;
      TFitResultPtr r=cr->Fit(f2,"SQ");

      Int_t status   = r;
      Double_t chi2  = r->Chi2()/r->Ndf();
      Double_t slope = r->Parameter(0);
      Double_t error = r->ParError(0);
      cout<<endl<<pnam[i]<<"  \tstatus: "<<status<<"   chi2/ndf: "<<chi2<<
	"\tc*tau: "<<slope<<"+/-"<<error<<endl<<endl;
  }

  return;
}

Double_t myExp2(Double_t *v, Double_t *p) {
   Double_t pt=v[0];
   Double_t lt=v[1];

   Double_t ctau=p[1];
   Double_t mass=p[2];
   Double_t ct=mass*lt/pt;
   if ((lt < 2) || (ct > 2.5*ctau)) {
      TF1::RejectPoint();
      return 0;
   }

   Int_t i=pt/0.1 + 1 + 1;
   return p[i]*TMath::Exp(-ct/p[0]);
}

void Correct(TH1 *rw, const TH1 *as, const TH1 *mc) {
  TH1 *eff = (TH1*)as->Clone();
  eff->Divide(as,mc,1,1,"b");
  rw->Divide(eff);
  delete eff;
}

void Normalise(Double_t br, Double_t yw, Double_t bw, Double_t ne, TH1 *pt) {
   pt->Scale(1/br);    // branching ratio
   pt->Scale(1/yw);    // rapidity window
   pt->Scale(1/bw);    // bin width
   pt->Scale(1/ne);    // number of events
}
<commit_msg>Optional possibility to correct in 1D and a more stable calculation of the global FD correction<commit_after>#if !defined(__CINT__) || defined(__MAKECINT__)
   #include <TMath.h>
   #include <TROOT.h>
   #include <Riostream.h>
   #include <TCanvas.h>

   #include <TString.h>

   #include <TFile.h>
   #include <TList.h>
   #include <TH1F.h>
   #include <TH1D.h>
   #include <TF2.h>
   #include <TFitResult.h>
   #include <TFitResultPtr.h>
   #include <TH2F.h>
   #include <TH3F.h>
#endif

extern TROOT *gROOT;
//extern TFile *gFile;

void PostProcessCTau(Int_t corrType=1) {
  //-----------------------------------------------------------------
  // PostProcess the histograms produced by AliAnalysisTaskCTau*
  // Produce:
  // -- V0 particle spectra (not corrected for the feeddown)
  // -- Estimated (anti)Lambda feed-down spectra (from Xi's only)
  // -- V0 particle c*tau distributions (feeddown corrected)
  //-----------------------------------------------------------------

  TString name("cTau_0090"); //centrality
  //TString name("cTau_0005"); //centrality
  //TString name("cTau_2040"); //centrality
  //TString name("cTau_4060"); //centrality
  //TString name("cTau_6080"); //centrality
  //TString name("cTau_8090"); //centrality

  //+++++ Real data histograms
  TFile::Open("new/real/all/" + name + ".root");
  TList *list=(TList *)gFile->Get(name);

  TH2F *fK0sSi=(TH2F*)list->FindObject("fK0sSi");fK0sSi->Sumw2();
  TH2F *fLambdaSi=(TH2F*)list->FindObject("fLambdaSi");fLambdaSi->Sumw2();
  TH1F *fXiSiP=(TH1F*)list->FindObject("fXiSiP"); fXiSiP->Sumw2();
  TH2F *fLambdaBarSi=
        (TH2F*)list->FindObject("fLambdaBarSi"); fLambdaBarSi->Sumw2();
  TH1F *fXiBarSiP=(TH1F*)list->FindObject("fXiBarSiP"); fXiBarSiP->Sumw2();

  TH1F *fMult=(TH1F*)list->FindObject("fMult");
  //+++++

  const Double_t nEvents=fMult->Integral();

  //+++++ MC histograms
  TFile::Open("new/mc_b_plus/all/" + name + "_mc.root");
  TList *listmc=(TList *)gFile->Get(name+"_mc");

  TH2F *fK0sMC=(TH2F*)listmc->FindObject("fK0sMC"); fK0sMC->Sumw2();
  TH2F *fK0sAs=(TH2F*)listmc->FindObject("fK0sAs"); fK0sAs->Sumw2();

  TH2F *fLambdaMC=(TH2F*)listmc->FindObject("fLambdaMC"); fLambdaMC->Sumw2();
  TH2F *fLambdaAs=(TH2F*)listmc->FindObject("fLambdaAs"); fLambdaAs->Sumw2();

  TH2F *fLambdaBarMC=
        (TH2F*)listmc->FindObject("fLambdaBarMC"); fLambdaBarMC->Sumw2();
  TH2F *fLambdaBarAs=
        (TH2F*)listmc->FindObject("fLambdaBarAs"); fLambdaBarAs->Sumw2();

  TH1F *fXiSiPMC=(TH1F*)listmc->FindObject("fXiSiP");fXiSiPMC->Sumw2();
  TH3F *fLambdaFromXi=(TH3F*)listmc->FindObject("fLambdaFromXi");
  fLambdaFromXi->Sumw2();

  TH1F *fXiBarSiPMC=(TH1F*)listmc->FindObject("fXiBarSiP");fXiBarSiPMC->Sumw2();
  TH3F *fLambdaBarFromXiBar=(TH3F*)listmc->FindObject("fLambdaBarFromXiBar");
  fLambdaBarFromXiBar->Sumw2();
  //+++++

  Double_t myExp2(Double_t *v, Double_t *p);
  void Correct(TH1 *rw, const TH1 *as, const TH1 *mc);
  void Normalise(Double_t, Double_t, Double_t, Double_t, TH1 *h);

  const Int_t    iMax=3;   // Number of V0 particles
  const TString  pnam[]={"K^{0}_{S}", "#Lambda", "#bar{#Lambda}"};
  const Double_t brch[]={0.69, 0.64, 0.64};
  const Double_t mass[]={0.4977, 1.115, 1.115};
  const Double_t ctau[]={2.68, 7.89, 7.89};
  const Double_t yWin=2*0.5; // rapidity window
  const Double_t wbx=fK0sSi->GetBinWidth(3);
  const Int_t nbx=fLambdaFromXi->GetNbinsX();
  const Int_t nby=fLambdaFromXi->GetNbinsY();
  const Int_t nbz=fLambdaFromXi->GetNbinsZ();

  const TH2 *in[]={
    fK0sSi,fK0sAs,fK0sMC,
    fLambdaSi,fLambdaAs,fLambdaMC,
    fLambdaBarSi,fLambdaBarAs,fLambdaBarMC
  };
  TH1 *fd[]={
    0,0,0,
    fLambdaFromXi,fXiSiP,fXiSiPMC,
    fLambdaBarFromXiBar,fXiBarSiP,fXiBarSiPMC
  };

  for (Int_t i=0; i<iMax; i++) {
      const TH2 *ptRw=in[3*i+0], *ptAs=in[3*i+1], *ptMc=in[3*i+2];
      TH1 *ptAsPx=0, *ptMcPx=0, *pt=0;

      TH2 *cr=(TH2*)ptRw->Clone();
      Correct(cr, ptAs, ptMc);

      //++++ pT spectrum
      if (corrType==1) {
         pt    =ptRw->ProjectionX("_px",0,-1,"e");
         ptAsPx=ptAs->ProjectionX("_px",0,-1,"e");
         ptMcPx=ptMc->ProjectionX("_px",0,-1,"e");
         Correct(pt, ptAsPx, ptMcPx);
      } else {
         pt=cr->ProjectionX("_px",0,-1,"e");
      }
      TString ptName = pnam[i] + " p_{T} spectrum";
      pt->SetTitle(ptName.Data());
      Normalise(brch[i], yWin, wbx, nEvents, pt);

      new TCanvas;
      pt->Draw();

      if (i>0) {
      //++++ feeddown
	 TH3 *fd3=(TH3*)fd[3*i + 0];
	 TH1 *rl =(TH1*)fd[3*i + 1];
	 TH1 *mc =(TH1*)fd[3*i + 2];

         Double_t nLambdaFromXiMc=fd3->Integral();
         Double_t nXiMc=mc->Integral();
         Double_t nXiRl=rl->Integral();
         Double_t f=(nLambdaFromXiMc*nXiRl/nXiMc)/ptRw->Integral();
         Double_t ef=f*TMath::Sqrt(nXiMc)/nXiMc;
         cout<<endl<<"Global FD correction: "<<f<<"+/-"<<ef<<endl;

         rl->Divide(mc);

         for (Int_t k=1; k<=nbx; k++) {
             for (Int_t l=1; l<=nby; l++) {
                 for (Int_t m=1; m<=nbz; m++) {
                     Float_t c=fd3->GetBinContent(k,l,m);
                     c *= rl->GetBinContent(m);
                     fd3->SetBinContent(k,l,m,c);
		 }
	     }
	 }

         TH2 *fd2=(TH2*)fd3->Project3D("yxe");
         Correct(fd2, ptAs, ptMc);

         TH1 *fd1=0;
         if (corrType==1) {
            fd1=(TH1*)fd3->Project3D("x");
            Correct(fd1, ptAsPx, ptMcPx);
         } else {
            fd1=fd2->ProjectionX("_px",0,-1,"e");
	 }
         Normalise(brch[i], yWin, wbx, nEvents, fd1);

         //new TCanvas();
         fd1->Draw("same");

         cr->Add(fd2,-1);
      }
      //continue;

      //++++ c*tau
      TF2 *f2=new TF2("myexpo2",myExp2,0.,10.,0.,100.,1+1+1+nbx);
      f2->SetParameter(0, ctau[i]);
      f2->FixParameter(1, ctau[i]);
      f2->FixParameter(2, mass[i]);
      for (Int_t p=1+1+1; p<=nbx+1+1; p++)
          f2->SetParameter(p,fK0sSi->GetBinContent(p,1));

      new TCanvas;
      TFitResultPtr r=cr->Fit(f2,"SQ");

      Int_t status   = r;
      Double_t chi2  = r->Chi2()/r->Ndf();
      Double_t slope = r->Parameter(0);
      Double_t error = r->ParError(0);
      cout<<endl<<pnam[i]<<"  \tstatus: "<<status<<"   chi2/ndf: "<<chi2<<
	"\tc*tau: "<<slope<<"+/-"<<error<<endl<<endl;
  }

  return;
}

Double_t myExp2(Double_t *v, Double_t *p) {
   Double_t pt=v[0];
   Double_t lt=v[1];

   Double_t ctau=p[1];
   Double_t mass=p[2];
   Double_t ct=mass*lt/pt;
   if ((lt < 2) || (ct > 2.5*ctau)) {
      TF1::RejectPoint();
      return 0;
   }

   Int_t i=pt/0.1 + 1 + 1;
   return p[i]*TMath::Exp(-ct/p[0]);
}

void Correct(TH1 *rw, const TH1 *as, const TH1 *mc) {
  TH1 *eff = (TH1*)as->Clone();
  eff->Divide(as,mc,1,1,"b");
  rw->Divide(eff);
  delete eff;
}

void Normalise(Double_t br, Double_t yw, Double_t bw, Double_t ne, TH1 *pt) {
   pt->Scale(1/br);    // branching ratio
   pt->Scale(1/yw);    // rapidity window
   pt->Scale(1/bw);    // bin width
   pt->Scale(1/ne);    // number of events
}
<|endoftext|>


<commit_before>/* This software is released under the BSD License.
 |
 | Copyright (c) 2015, Kevin P. Barry [ta0kira@gmail.com]
 | 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  Locking Container Project  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 COPYRIGHT  OWNER  OR 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.
 +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~*/

#ifndef authorization_hpp
#define authorization_hpp

#include <assert.h>

/*! \class lock_auth_base
 *  \brief Base class for lock authorization classes.
 *
 * @see lock_auth
 */

class lock_base;

class lock_auth_base {
public:
  typedef long count_type;
  typedef std::shared_ptr <lock_auth_base> auth_type;

  virtual count_type reading_count() const { return 0; }
  virtual count_type writing_count() const { return 0; }

  /*! Attempt to predict if a read authorization would be granted.*/
  inline bool guess_read_allowed(bool lock_out = true, bool in_use = true) const {
    return this->test_auth(true, lock_out, in_use);
  }

  /*! Attempt to predict if a write authorization would be granted.*/
  inline bool guess_write_allowed(bool lock_out = true, bool in_use = true) const {
    return this->test_auth(false, lock_out, in_use);
  }

  virtual inline ~lock_auth_base() {}

private:
  friend class lock_base;

  /*! Obtain lock authorization.*/
  virtual bool register_auth(bool read, bool lock_out, bool in_use) = 0;

  /*! Test lock authorization.*/
  virtual bool test_auth(bool read, bool lock_out, bool in_use) const = 0;

  /*! Release lock authorization.*/
  virtual void release_auth(bool read) = 0;
};

/*! \class lock_auth
 *  \brief Lock authorization object.
 *
 * @see locking_container::auth_type
 * @see locking_container::get_new_auth
 * @see locking_container::get_auth
 * @see locking_container::get_auth_const
 *
 * This class is used by \ref locking_container to prevent deadlocks. To prevent
 * deadlocks, create one \ref lock_auth instance per thread, and pass it to the
 * \ref locking_container when getting a proxy object. This will prevent the
 * thread from obtaining an new incompatible lock type when it already holds a
 * lock.
 */

template <class> class lock_auth;

/*! \class lock_auth <rw_lock>
 *
 * This auth. type allows the caller to hold multiple read locks, or a single
 * write lock, but not both. Note that if another thread is waiting for a write
 * lock on a w_lock or rw_lock container and the caller already has a read lock
 * then the lock will be rejected. Three exceptions to these rules are: 1) if
 * the container to be locked currently has no other locks; 2) if the call isn't
 * blocking and it's for a write lock; 3) if the container uses a rw_lock, the
 * caller has a write lock on the container, and the caller is requesting a read
 * lock. (The 3rd case is what allows multi-lock to work.)
 */

class rw_lock;

template <>
class lock_auth <rw_lock> : public lock_auth_base {
public:
  using lock_auth_base::count_type;

  lock_auth() : reading(0), writing(0) {}

  count_type reading_count() const { return reading; }
  count_type writing_count() const { return writing; }

  ~lock_auth() {
    //NOTE: this can't be in '~lock_auth_base'!
    assert(!this->reading_count() && !this->writing_count());
  }

private:
  lock_auth(const lock_auth&);
  lock_auth &operator = (const lock_auth&);

  bool register_auth(bool read, bool lock_out, bool in_use) {
    if (!this->test_auth(read, lock_out, in_use)) return false;
    if (read) {
      ++reading;
      assert(reading > 0);
    } else {
      ++writing;
      assert(writing > 0);
    }
    return true;
  }

  bool test_auth(bool read, bool lock_out, bool in_use) const {
    if (writing && in_use)                return false;
    if (reading && !read && in_use)       return false;
    if ((reading || writing) && lock_out) return false;
    return true;
  }

  void release_auth(bool read) {
    if (read) {
      //NOTE: don't check 'writing' because there are a few exceptions!
      assert(reading > 0);
      --reading;
    } else {
      //NOTE: don't check 'reading' because there are a few exceptions!
      assert(writing > 0);
      --writing;
    }
  }

  count_type reading, writing;
};

/*! \class lock_auth <r_lock>
 *
 * This auth. type allows the caller to hold multiple read locks, but no write
 * locks. Note that if another thread is waiting for a write lock on the
 * container and the caller already has a read lock then the lock will be
 * rejected. Use this auth. type if you want to ensure that a thread doesn't
 * obtain a write lock on any container. Using this auth. type with containers
 * that use any combination of r_lock and rw_lock should result in the same
 * behavior as using it with only r_lock containers.
 */

class r_lock;

template <>
class lock_auth <r_lock> : public lock_auth_base {
public:
  using lock_auth_base::count_type;

  lock_auth() : reading(0) {}

  count_type reading_count() const { return reading; }

  ~lock_auth() {
    //NOTE: this can't be in '~lock_auth_base'!
    //NOTE: no point checking 'writing_count', since it's overrides will be ignored here
    assert(!this->reading_count());
  }

private:
  bool register_auth(bool read, bool lock_out, bool in_use) {
    if (!this->test_auth(read, lock_out, in_use)) return false;
    ++reading;
    assert(reading > 0);
    return true;
  }

  bool test_auth(bool read, bool lock_out, bool /*in_use*/) const {
    if (!read)               return false;
    if (reading && lock_out) return false;
    return true;
  }

  void release_auth(bool read) {
    assert(read);
    assert(reading > 0);
    --reading;
  }

  count_type reading;
};

/*! \class lock_auth <w_lock>
 *
 * This auth. type allows the caller to hold no more than one lock at a time,
 * regardless of lock type. An exception to this behavior is if the container to
 * be locked currently has no other locks. Use this auth. type if you are only
 * using containers that use w_lock, or if you want to disallow multiple read
 * locks on containers that use rw_lock or r_lock. Using this auth. type with
 * containers that use any combination of w_lock and rw_lock should result in
 * the same behavior as using it with only w_lock containers.
 */

class w_lock;

template <>
class lock_auth <w_lock> : public lock_auth_base {
public:
  using lock_auth_base::count_type;

  lock_auth() : writing(0) {}

  count_type writing_count() const { return writing; }

  ~lock_auth() {
    //NOTE: this can't be in '~lock_auth_base'!
    //NOTE: no point checking 'reading_count', since it's overrides will be ignored here
    assert(!this->writing_count());
  }

private:
  lock_auth(const lock_auth&);
  lock_auth &operator = (const lock_auth&);

  bool register_auth(bool read, bool lock_out, bool in_use) {
    if (!this->test_auth(read, lock_out, in_use)) return false;
    ++writing;
    assert(writing > 0);
    return true;
  }

  bool test_auth(bool /*read*/, bool /*lock_out*/, bool in_use) const {
    return !writing || !in_use;
  }

  void release_auth(bool /*read*/) {
    assert(writing > 0);
    --writing;
  }

  count_type writing;
};

/*! \class lock_auth <dumb_lock>
 *
 * This auth. type only allows the caller to hold one lock at a time. Unlike
 * lock_auth <w_lock>, it doesn't matter if the container is in use. (One caveat
 * to this is that r_lock doesn't actually lock; therefore, this auth. can
 * potentially hold multiple locks on r_lock containers at one time.) This is
 * useful if you want to ensure that the caller can only hold a single lock at
 * any given time. This auth. type will not work with multi-locking.
 */

class dumb_lock;

template <>
class lock_auth <dumb_lock> : public lock_auth_base {
public:
  using lock_auth_base::count_type;

  lock_auth() : writing(false) {}

  count_type writing_count() const { return writing? 1 : 0; }

  ~lock_auth() {
    //NOTE: this can't be in '~lock_auth_base'!
    //NOTE: no point checking 'reading_count', since it's overrides will be ignored here
    assert(!this->writing_count());
  }

private:
  lock_auth(const lock_auth&);
  lock_auth &operator = (const lock_auth&);

  bool register_auth(bool read, bool lock_out, bool in_use) {
    if (!this->test_auth(read, lock_out, in_use)) return false;
    writing = true;
    return true;
  }

  bool test_auth(bool /*read*/, bool /*lock_out*/, bool /*in_use*/) const {
    return !writing;
  }

  void release_auth(bool /*read*/) {
    assert(writing);
    writing = false;
  }

  bool writing;
};

/*! \class lock_auth <broken_lock>
 *
 * This auth. type doesn't allow the caller to obtain any locks.
 */

struct broken_lock;

template <>
class lock_auth <broken_lock> : public lock_auth_base {
public:
  using lock_auth_base::count_type;

private:
  bool register_auth(bool /*read*/, bool /*lock_out*/, bool /*in_use*/)   { return false; }
  bool test_auth(bool /*read*/, bool /*lock_out*/, bool /*in_use*/) const { return false; }
  void release_auth(bool /*read*/)                                        { assert(false); }
};

#endif //authorization_hpp
<commit_msg>revised definitions of lock_auth to allow inheritance<commit_after>/* This software is released under the BSD License.
 |
 | Copyright (c) 2015, Kevin P. Barry [ta0kira@gmail.com]
 | 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  Locking Container Project  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 COPYRIGHT  OWNER  OR 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.
 +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~*/

#ifndef authorization_hpp
#define authorization_hpp

#include <assert.h>

/*! \class lock_auth_base
 *  \brief Base class for lock authorization classes.
 *
 * @see lock_auth
 */

class lock_base;

class lock_auth_base {
public:
  typedef long count_type;
  typedef std::shared_ptr <lock_auth_base> auth_type;

  virtual count_type reading_count() const { return 0; }
  virtual count_type writing_count() const { return 0; }

  /*! Attempt to predict if a read authorization would be granted.*/
  inline bool guess_read_allowed(bool lock_out = true, bool in_use = true) const {
    return this->test_auth(true, lock_out, in_use);
  }

  /*! Attempt to predict if a write authorization would be granted.*/
  inline bool guess_write_allowed(bool lock_out = true, bool in_use = true) const {
    return this->test_auth(false, lock_out, in_use);
  }

  virtual inline ~lock_auth_base() {}

private:
  friend class lock_base;

  /*! Obtain lock authorization.*/
  virtual bool register_auth(bool read, bool lock_out, bool in_use) = 0;

  /*! Test lock authorization.*/
  virtual bool test_auth(bool read, bool lock_out, bool in_use) const = 0;

  /*! Release lock authorization.*/
  virtual void release_auth(bool read) = 0;
};

/*! \class lock_auth
 *  \brief Lock authorization object.
 *
 * @see locking_container::auth_type
 * @see locking_container::get_new_auth
 * @see locking_container::get_auth
 * @see locking_container::get_auth_const
 *
 * This class is used by \ref locking_container to prevent deadlocks. To prevent
 * deadlocks, create one \ref lock_auth instance per thread, and pass it to the
 * \ref locking_container when getting a proxy object. This will prevent the
 * thread from obtaining an new incompatible lock type when it already holds a
 * lock.
 */

template <class> class lock_auth;

/*! \class lock_auth_rw_lock
 *
 * This auth. type allows the caller to hold multiple read locks, or a single
 * write lock, but not both. Note that if another thread is waiting for a write
 * lock on a w_lock or rw_lock container and the caller already has a read lock
 * then the lock will be rejected. Three exceptions to these rules are: 1) if
 * the container to be locked currently has no other locks; 2) if the call isn't
 * blocking and it's for a write lock; 3) if the container uses a rw_lock, the
 * caller has a write lock on the container, and the caller is requesting a read
 * lock. (The 3rd case is what allows multi-lock to work.)
 */

class lock_auth_rw_lock : public lock_auth_base {
public:
  using lock_auth_base::count_type;

  lock_auth_rw_lock() : reading(0), writing(0) {}

  count_type reading_count() const { return reading; }
  count_type writing_count() const { return writing; }

  ~lock_auth_rw_lock() {
    //NOTE: this can't be in '~lock_auth_base'!
    assert(!this->reading_count() && !this->writing_count());
  }

private:
  lock_auth_rw_lock(const lock_auth_rw_lock&);
  lock_auth_rw_lock &operator = (const lock_auth_rw_lock&);

  bool register_auth(bool read, bool lock_out, bool in_use) {
    if (!this->test_auth(read, lock_out, in_use)) return false;
    if (read) {
      ++reading;
      assert(reading > 0);
    } else {
      ++writing;
      assert(writing > 0);
    }
    return true;
  }

  bool test_auth(bool read, bool lock_out, bool in_use) const {
    if (writing && in_use)                return false;
    if (reading && !read && in_use)       return false;
    if ((reading || writing) && lock_out) return false;
    return true;
  }

  void release_auth(bool read) {
    if (read) {
      //NOTE: don't check 'writing' because there are a few exceptions!
      assert(reading > 0);
      --reading;
    } else {
      //NOTE: don't check 'reading' because there are a few exceptions!
      assert(writing > 0);
      --writing;
    }
  }

  count_type reading, writing;
};

class rw_lock;

template <>
class lock_auth <rw_lock> : public lock_auth_rw_lock {};

/*! \class lock_auth_r_lock
 *
 * This auth. type allows the caller to hold multiple read locks, but no write
 * locks. Note that if another thread is waiting for a write lock on the
 * container and the caller already has a read lock then the lock will be
 * rejected. Use this auth. type if you want to ensure that a thread doesn't
 * obtain a write lock on any container. Using this auth. type with containers
 * that use any combination of r_lock and rw_lock should result in the same
 * behavior as using it with only r_lock containers.
 */

class lock_auth_r_lock : public lock_auth_base {
public:
  using lock_auth_base::count_type;

  lock_auth_r_lock() : reading(0) {}

  count_type reading_count() const { return reading; }

  ~lock_auth_r_lock() {
    //NOTE: this can't be in '~lock_auth_base'!
    //NOTE: no point checking 'writing_count', since it's overrides will be ignored here
    assert(!this->reading_count());
  }

private:
  lock_auth_r_lock(const lock_auth_r_lock&);
  lock_auth_r_lock &operator = (const lock_auth_r_lock&);

  bool register_auth(bool read, bool lock_out, bool in_use) {
    if (!this->test_auth(read, lock_out, in_use)) return false;
    ++reading;
    assert(reading > 0);
    return true;
  }

  bool test_auth(bool read, bool lock_out, bool /*in_use*/) const {
    if (!read)               return false;
    if (reading && lock_out) return false;
    return true;
  }

  void release_auth(bool read) {
    assert(read);
    assert(reading > 0);
    --reading;
  }

  count_type reading;
};

class r_lock;

template <>
class lock_auth <r_lock> : public lock_auth_r_lock {};

/*! \class lock_auth_w_lock
 *
 * This auth. type allows the caller to hold no more than one lock at a time,
 * regardless of lock type. An exception to this behavior is if the container to
 * be locked currently has no other locks. Use this auth. type if you are only
 * using containers that use w_lock, or if you want to disallow multiple read
 * locks on containers that use rw_lock or r_lock. Using this auth. type with
 * containers that use any combination of w_lock and rw_lock should result in
 * the same behavior as using it with only w_lock containers.
 */

class lock_auth_w_lock : public lock_auth_base {
public:
  using lock_auth_base::count_type;

  lock_auth_w_lock() : writing(0) {}

  count_type writing_count() const { return writing; }

  ~lock_auth_w_lock() {
    //NOTE: this can't be in '~lock_auth_base'!
    //NOTE: no point checking 'reading_count', since it's overrides will be ignored here
    assert(!this->writing_count());
  }

private:
  lock_auth_w_lock(const lock_auth_w_lock&);
  lock_auth_w_lock &operator = (const lock_auth_w_lock&);

  bool register_auth(bool read, bool lock_out, bool in_use) {
    if (!this->test_auth(read, lock_out, in_use)) return false;
    ++writing;
    assert(writing > 0);
    return true;
  }

  bool test_auth(bool /*read*/, bool /*lock_out*/, bool in_use) const {
    return !writing || !in_use;
  }

  void release_auth(bool /*read*/) {
    assert(writing > 0);
    --writing;
  }

  count_type writing;
};

class w_lock;

template <>
class lock_auth <w_lock> : public lock_auth_w_lock {};

/*! \class lock_auth_dumb_lock
 *
 * This auth. type only allows the caller to hold one lock at a time. Unlike
 * lock_auth_w_lock, it doesn't matter if the container is in use. (One caveat
 * to this is that r_lock doesn't actually lock; therefore, this auth. can
 * potentially hold multiple locks on r_lock containers at one time.) This is
 * useful if you want to ensure that the caller can only hold a single lock at
 * any given time. This auth. type will not work with multi-locking.
 */

class lock_auth_dumb_lock : public lock_auth_base {
public:
  using lock_auth_base::count_type;

  lock_auth_dumb_lock() : writing(false) {}

  count_type writing_count() const { return writing? 1 : 0; }

  ~lock_auth_dumb_lock() {
    //NOTE: this can't be in '~lock_auth_base'!
    //NOTE: no point checking 'reading_count', since it's overrides will be ignored here
    assert(!this->writing_count());
  }

private:
  lock_auth_dumb_lock(const lock_auth_dumb_lock&);
  lock_auth_dumb_lock &operator = (const lock_auth_dumb_lock&);

  bool register_auth(bool read, bool lock_out, bool in_use) {
    if (!this->test_auth(read, lock_out, in_use)) return false;
    writing = true;
    return true;
  }

  bool test_auth(bool /*read*/, bool /*lock_out*/, bool /*in_use*/) const {
    return !writing;
  }

  void release_auth(bool /*read*/) {
    assert(writing);
    writing = false;
  }

  bool writing;
};

class dumb_lock;

template <>
class lock_auth <dumb_lock> : public lock_auth_dumb_lock {};

/*! \class lock_auth_broken_lock
 *
 * This auth. type doesn't allow the caller to obtain any locks.
 */

class lock_auth_broken_lock : public lock_auth_base {
public:
  using lock_auth_base::count_type;

private:
  bool register_auth(bool /*read*/, bool /*lock_out*/, bool /*in_use*/)   { return false; }
  bool test_auth(bool /*read*/, bool /*lock_out*/, bool /*in_use*/) const { return false; }
  void release_auth(bool /*read*/)                                        { assert(false); }
};

struct broken_lock;

template <>
class lock_auth <broken_lock> : public lock_auth_broken_lock {};

#endif //authorization_hpp
<|endoftext|>


<commit_before>/* BEGIN_COMMON_COPYRIGHT_HEADER
 * (c)LGPL2+
 *
 * LXDE-Qt - a lightweight, Qt based, desktop toolset
 * http://razor-qt.org
 *
 * Copyright: 2012 Christian Surlykke
 *
 * This program or library is free software; you can redistribute it
 * and/or modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library 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 for more details.

 * You should have received a copy of the GNU Lesser General
 * Public License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
 * Boston, MA 02110-1301 USA
 *
 * END_COMMON_COPYRIGHT_HEADER */

#include "loginform.h"
#include "ui_loginform.h"

#include <QtCore/QDebug>
#include <QtGui/QCompleter>
#include <QtCore/QAbstractListModel>
#include <QtCore/QModelIndex>
#include <QtCore/QFile>
#include <QtCore/QTextStream>
#include <QtCore/QStringList>
#include <QtGui/QPixmap>
#include <QtGui/QMessageBox>
#include <QtGui/QMenu>
#include <QtCore/QProcess>
#include <lxqt/lxqtsettings.h>

LoginForm::LoginForm(QWidget *parent) :
    QWidget(parent),
        ui(new Ui::LoginForm),
        m_Greeter(),
        m_LoginData(&m_Greeter),
        m_powerManager(this, true),
        m_otherUserComboIndex(-1)
{
    if (!m_Greeter.connectSync())
    {
        close();
    }

    ui->setupUi(this);

    setupAppearence();
    fillUserAndSessionCombos();
    setupConnections();
    initializeControls();
    ui->formFrame->adjustSize();
    adjustSize();
}

LoginForm::~LoginForm()
{
    delete ui;
}

void LoginForm::setFocus(Qt::FocusReason reason)
{
    if (ui->userCombo->currentIndex() == -1)
    {
        ui->userCombo->setFocus(reason);
    }
    else if (ui->userCombo->currentIndex() == m_otherUserComboIndex)
    {
        ui->otherUserInput->setFocus(reason);
    }
    else
    {
        ui->passwordInput->setFocus(reason);
    }
}

void LoginForm::setupAppearence()
{
    QPixmap icon(":/resources/helix_1120.png");
    ui->iconLabel->setPixmap(icon.scaled(ui->iconLabel->size(), Qt::KeepAspectRatio, Qt::SmoothTransformation));
    ui->hostnameLabel->setText(m_Greeter.hostname());
}

void LoginForm::fillUserAndSessionCombos()
{
    ui->userCombo->clear();
    for (int i = 0; i < m_LoginData.numberOfUsers(); i++)
    {
        qDebug() << "Adding user" << m_LoginData.userFullName(i);
        ui->userCombo->addItem(m_LoginData.userFullName(i));
    }
    ui->userCombo->addItem(tr("other..."));
    m_otherUserComboIndex = ui->userCombo->count() - 1;

    ui->sessionCombo->clear();
    for (int i = 0; i < m_LoginData.numberOfSessions(); i++)
    {
        qDebug() << "Adding session" << m_LoginData.sessionFullName(i);
        ui->sessionCombo->addItem(m_LoginData.sessionFullName(i));
    }
}

void LoginForm::setupConnections()
{
    ui->userCombo->setCurrentIndex(-1);
    connect(ui->userCombo, SIGNAL(currentIndexChanged(int)), this, SLOT(userComboCurrentIndexChanged()));
    connect(ui->otherUserInput, SIGNAL(editingFinished()), this, SLOT(otherUserEditingFinished()));
    connect(ui->loginButton, SIGNAL(clicked(bool)), this, SLOT(loginClicked()));
    connect(ui->leaveButton, SIGNAL(clicked()), this, SLOT(leaveClicked()));

    connect(&m_Greeter, SIGNAL(showPrompt(QString,QLightDM::Greeter::PromptType)),
            this,       SLOT(onPrompt(QString,QLightDM::Greeter::PromptType)));

    connect(&m_Greeter, SIGNAL(authenticationComplete()), this, SLOT(authenticationComplete()));
}

void LoginForm::initializeControls()
{
    QMenu* leaveMenu = new QMenu(this);
    ui->leaveButton->setMenu(leaveMenu);

    foreach (QAction *action, m_powerManager.availableActions())
    {
        leaveMenu->addAction(action);
    }

    qDebug() << "showManualLoginHint:" << m_Greeter.showManualLoginHint();

    ui->sessionCombo->setCurrentIndex(m_LoginData.suggestedSession());

    if (m_Greeter.showManualLoginHint())
    {
        ui->userCombo->setCurrentIndex(m_otherUserComboIndex);
    }
    else
    {
        ui->userCombo->setCurrentIndex(m_LoginData.suggestedUser());
    }

    ui->userCombo->setVisible(! m_Greeter.hideUsersHint());

    ui->otherUserInput->setVisible(ui->userCombo->currentIndex() == m_otherUserComboIndex);
    ui->otherUserInput->clear();
    QStringList completions;
    for (int i = 0; i < m_LoginData.numberOfUsers(); i++)
    {
        completions << m_LoginData.userName(i);
    }
    QCompleter *completer = new QCompleter(completions, ui->otherUserInput);
    completer->setCompletionMode(QCompleter::InlineCompletion);
    ui->otherUserInput->setCompleter(completer);

    ui->passwordInput->setEnabled(false);
    ui->passwordInput->clear();
}

void LoginForm::setSessionCombo(int session_index)
{
    if (0 <= session_index && session_index < ui->sessionCombo->count())
    {
        ui->sessionCombo->setCurrentIndex(session_index);
    }
}

void LoginForm::userComboCurrentIndexChanged()
{

    qDebug() << "userComboCurrentIndexChanged:" << ui->userCombo->currentIndex();
    qDebug() << "setVisible...";
    if (ui->userCombo->currentIndex() == m_otherUserComboIndex)
    {
        ui->otherUserInput->show();
        ui->otherUserInput->setFocus();
    }
    else
    {
        ui->otherUserInput->hide();
        qDebug() << "Start authentication..";
        if (ui->userCombo->currentIndex() > -1)
        {
            setUser(m_LoginData.userName(ui->userCombo->currentIndex()));
        }
    }
    ui->formFrame->adjustSize();
    adjustSize();
}

void LoginForm::otherUserEditingFinished()
{
    if (ui->otherUserInput->text().isNull() || ui->otherUserInput->text().isEmpty())
    {
        return;
    }

    qDebug() << "Authenticating with otherUser...";

    setUser(ui->otherUserInput->text()) ;
}

void LoginForm::loginClicked()
{
    qDebug() << "loginClicked";
    m_Greeter.respond(ui->passwordInput->text().trimmed());
    ui->passwordInput->clear();
    ui->passwordInput->setEnabled(false);
}

void LoginForm::onPrompt(QString prompt, QLightDM::Greeter::PromptType promptType)
{
    qDebug() << "onPrompt";
    ui->passwordInput->setEnabled(true);
    ui->passwordInput->setFocus();
}

void LoginForm::setUser(QString user)
{
    m_user = user;
    setSessionCombo(m_LoginData.userSession(m_user));

    if (m_Greeter.inAuthentication())
    {
        m_Greeter.cancelAuthentication();
    }
    m_Greeter.authenticate(m_user);
}

void LoginForm::authenticationComplete()
{
    if (m_Greeter.isAuthenticated())
    {
        qDebug() << "Authenticated... starting session" << m_LoginData.sessionName(ui->sessionCombo->currentIndex());
        m_LoginData.setLastUser(m_user);
        m_Greeter.startSessionSync(m_LoginData.sessionName(ui->sessionCombo->currentIndex()));
        initializeControls(); // We should not arrive here...
    }
    else
    {
        qDebug() << "Not authenticated";
        ui->passwordInput->clear();
        setUser(m_user);
    }
}

<commit_msg>Remove connect to non-existing slot in loginform.cpp<commit_after>/* BEGIN_COMMON_COPYRIGHT_HEADER
 * (c)LGPL2+
 *
 * LXDE-Qt - a lightweight, Qt based, desktop toolset
 * http://razor-qt.org
 *
 * Copyright: 2012 Christian Surlykke
 *
 * This program or library is free software; you can redistribute it
 * and/or modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library 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 for more details.

 * You should have received a copy of the GNU Lesser General
 * Public License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
 * Boston, MA 02110-1301 USA
 *
 * END_COMMON_COPYRIGHT_HEADER */

#include "loginform.h"
#include "ui_loginform.h"

#include <QtCore/QDebug>
#include <QtGui/QCompleter>
#include <QtCore/QAbstractListModel>
#include <QtCore/QModelIndex>
#include <QtCore/QFile>
#include <QtCore/QTextStream>
#include <QtCore/QStringList>
#include <QtGui/QPixmap>
#include <QtGui/QMessageBox>
#include <QtGui/QMenu>
#include <QtCore/QProcess>
#include <lxqt/lxqtsettings.h>

LoginForm::LoginForm(QWidget *parent) :
    QWidget(parent),
        ui(new Ui::LoginForm),
        m_Greeter(),
        m_LoginData(&m_Greeter),
        m_powerManager(this, true),
        m_otherUserComboIndex(-1)
{
    if (!m_Greeter.connectSync())
    {
        close();
    }

    ui->setupUi(this);

    setupAppearence();
    fillUserAndSessionCombos();
    setupConnections();
    initializeControls();
    ui->formFrame->adjustSize();
    adjustSize();
}

LoginForm::~LoginForm()
{
    delete ui;
}

void LoginForm::setFocus(Qt::FocusReason reason)
{
    if (ui->userCombo->currentIndex() == -1)
    {
        ui->userCombo->setFocus(reason);
    }
    else if (ui->userCombo->currentIndex() == m_otherUserComboIndex)
    {
        ui->otherUserInput->setFocus(reason);
    }
    else
    {
        ui->passwordInput->setFocus(reason);
    }
}

void LoginForm::setupAppearence()
{
    QPixmap icon(":/resources/helix_1120.png");
    ui->iconLabel->setPixmap(icon.scaled(ui->iconLabel->size(), Qt::KeepAspectRatio, Qt::SmoothTransformation));
    ui->hostnameLabel->setText(m_Greeter.hostname());
}

void LoginForm::fillUserAndSessionCombos()
{
    ui->userCombo->clear();
    for (int i = 0; i < m_LoginData.numberOfUsers(); i++)
    {
        qDebug() << "Adding user" << m_LoginData.userFullName(i);
        ui->userCombo->addItem(m_LoginData.userFullName(i));
    }
    ui->userCombo->addItem(tr("other..."));
    m_otherUserComboIndex = ui->userCombo->count() - 1;

    ui->sessionCombo->clear();
    for (int i = 0; i < m_LoginData.numberOfSessions(); i++)
    {
        qDebug() << "Adding session" << m_LoginData.sessionFullName(i);
        ui->sessionCombo->addItem(m_LoginData.sessionFullName(i));
    }
}

void LoginForm::setupConnections()
{
    ui->userCombo->setCurrentIndex(-1);
    connect(ui->userCombo, SIGNAL(currentIndexChanged(int)), this, SLOT(userComboCurrentIndexChanged()));
    connect(ui->otherUserInput, SIGNAL(editingFinished()), this, SLOT(otherUserEditingFinished()));
    connect(ui->loginButton, SIGNAL(clicked(bool)), this, SLOT(loginClicked()));

    connect(&m_Greeter, SIGNAL(showPrompt(QString,QLightDM::Greeter::PromptType)),
            this,       SLOT(onPrompt(QString,QLightDM::Greeter::PromptType)));

    connect(&m_Greeter, SIGNAL(authenticationComplete()), this, SLOT(authenticationComplete()));
}

void LoginForm::initializeControls()
{
    QMenu* leaveMenu = new QMenu(this);
    ui->leaveButton->setMenu(leaveMenu);

    foreach (QAction *action, m_powerManager.availableActions())
    {
        leaveMenu->addAction(action);
    }

    qDebug() << "showManualLoginHint:" << m_Greeter.showManualLoginHint();

    ui->sessionCombo->setCurrentIndex(m_LoginData.suggestedSession());

    if (m_Greeter.showManualLoginHint())
    {
        ui->userCombo->setCurrentIndex(m_otherUserComboIndex);
    }
    else
    {
        ui->userCombo->setCurrentIndex(m_LoginData.suggestedUser());
    }

    ui->userCombo->setVisible(! m_Greeter.hideUsersHint());

    ui->otherUserInput->setVisible(ui->userCombo->currentIndex() == m_otherUserComboIndex);
    ui->otherUserInput->clear();
    QStringList completions;
    for (int i = 0; i < m_LoginData.numberOfUsers(); i++)
    {
        completions << m_LoginData.userName(i);
    }
    QCompleter *completer = new QCompleter(completions, ui->otherUserInput);
    completer->setCompletionMode(QCompleter::InlineCompletion);
    ui->otherUserInput->setCompleter(completer);

    ui->passwordInput->setEnabled(false);
    ui->passwordInput->clear();
}

void LoginForm::setSessionCombo(int session_index)
{
    if (0 <= session_index && session_index < ui->sessionCombo->count())
    {
        ui->sessionCombo->setCurrentIndex(session_index);
    }
}

void LoginForm::userComboCurrentIndexChanged()
{

    qDebug() << "userComboCurrentIndexChanged:" << ui->userCombo->currentIndex();
    qDebug() << "setVisible...";
    if (ui->userCombo->currentIndex() == m_otherUserComboIndex)
    {
        ui->otherUserInput->show();
        ui->otherUserInput->setFocus();
    }
    else
    {
        ui->otherUserInput->hide();
        qDebug() << "Start authentication..";
        if (ui->userCombo->currentIndex() > -1)
        {
            setUser(m_LoginData.userName(ui->userCombo->currentIndex()));
        }
    }
    ui->formFrame->adjustSize();
    adjustSize();
}

void LoginForm::otherUserEditingFinished()
{
    if (ui->otherUserInput->text().isNull() || ui->otherUserInput->text().isEmpty())
    {
        return;
    }

    qDebug() << "Authenticating with otherUser...";

    setUser(ui->otherUserInput->text()) ;
}

void LoginForm::loginClicked()
{
    qDebug() << "loginClicked";
    m_Greeter.respond(ui->passwordInput->text().trimmed());
    ui->passwordInput->clear();
    ui->passwordInput->setEnabled(false);
}

void LoginForm::onPrompt(QString prompt, QLightDM::Greeter::PromptType promptType)
{
    qDebug() << "onPrompt";
    ui->passwordInput->setEnabled(true);
    ui->passwordInput->setFocus();
}

void LoginForm::setUser(QString user)
{
    m_user = user;
    setSessionCombo(m_LoginData.userSession(m_user));

    if (m_Greeter.inAuthentication())
    {
        m_Greeter.cancelAuthentication();
    }
    m_Greeter.authenticate(m_user);
}

void LoginForm::authenticationComplete()
{
    if (m_Greeter.isAuthenticated())
    {
        qDebug() << "Authenticated... starting session" << m_LoginData.sessionName(ui->sessionCombo->currentIndex());
        m_LoginData.setLastUser(m_user);
        m_Greeter.startSessionSync(m_LoginData.sessionName(ui->sessionCombo->currentIndex()));
        initializeControls(); // We should not arrive here...
    }
    else
    {
        qDebug() << "Not authenticated";
        ui->passwordInput->clear();
        setUser(m_user);
    }
}

<|endoftext|>


<commit_before>#include "localsocket.h"

#include <errno.h>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include "sys/uio.h"
#include <sys/un.h>
#include <unistd.h>

#include <cassert>
#include <climits>
#include <cstdlib>
#include <cstring>

// HACK, put this somewhere else (get the value from original d-bus? or is it infinite?)
static const int maxFds = 12;

using namespace std;

LocalSocket::LocalSocket(int fd)
   : m_fd(fd)
{}

LocalSocket::LocalSocket(const string &socketFilePath)
   : m_fd(-1)
{
    const int fd = socket(PF_UNIX, SOCK_STREAM, 0);
    if (fd < 0) {
        return;
    }
    // don't let forks inherit the file descriptor - that can cause confusion...
    fcntl(fd, F_SETFD, FD_CLOEXEC);

    struct sockaddr_un addr;
    addr.sun_family = PF_UNIX;
    bool ok = socketFilePath.length() < sizeof(addr.sun_path);
    if (ok) {
        memcpy(addr.sun_path, socketFilePath.c_str(), socketFilePath.length());
    }

    ok = ok && (connect(fd, (struct sockaddr *)&addr, sizeof(sa_family_t) + socketFilePath.length()) == 0);

    if (ok) {
        m_fd = fd;
    } else {
        ::close(fd);
    }
}

LocalSocket::~LocalSocket()
{
    close();
}

void LocalSocket::close()
{
    setEventDispatcher(0);
    if (m_fd >= 0) {
        ::close(m_fd);
    }
    m_fd = -1;
}

int LocalSocket::write(array a)
{
    if (m_fd < 0) {
        return 0; // TODO -1?
    }

    // sendmsg  boilerplate
    struct msghdr send_msg; // C convention makes them stand out (too many similar variables...)
    struct iovec iov;

    send_msg.msg_name = 0;
    send_msg.msg_namelen = 0;
    send_msg.msg_flags = 0;
    send_msg.msg_iov = &iov;
    send_msg.msg_iovlen = 1;

    iov.iov_base = a.begin;
    iov.iov_len = a.length;

    // we can only send a fixed number of fds anyway due to the non-flexible size of the control message
    // receive buffer, so we set an arbitrary limit.
    const int numFds = 0; // TODO - how many should we get? do we need to know?
    assert(numFds <= maxFds);

    char cmsgBuf[CMSG_SPACE(sizeof(int) * maxFds)];

    if (numFds) {
        // fill in a control message
        send_msg.msg_control = cmsgBuf;
        send_msg.msg_controllen = CMSG_SPACE(sizeof(int) * numFds);

        struct cmsghdr *c_msg = CMSG_FIRSTHDR(&send_msg);
        c_msg->cmsg_len = CMSG_LEN(sizeof(int) * numFds);
        c_msg->cmsg_level = SOL_SOCKET;
        c_msg->cmsg_type = SCM_RIGHTS;

        // set the control data to pass - this is why we don't use the simpler write()
        for (int i = 0; i < numFds; i++) {
            // TODO
            // reinterpret_cast<int *>(CMSG_DATA(c_msg))[i] = message.m_fileDescriptors[i];
        }
    } else {
        // no file descriptor to send, no control message
        send_msg.msg_control = 0;
        send_msg.msg_controllen = 0;
    }

    while (iov.iov_len > 0) {
        int nbytes = sendmsg(m_fd, &send_msg, 0);
        if (nbytes < 0) {
            if (errno == EINTR) {
                continue;
            } else {
                close();
                return false;
            }
        }

        iov.iov_base = static_cast<char *>(iov.iov_base) + nbytes;
        iov.iov_len -= nbytes;
        // sendmsg() should always send the number of bytes asked for or block, so...
        if (nbytes != 0) {
            assert(iov.iov_len == 0);
        }
    }

    return a.length - iov.iov_len; // "- iov.iov_len" is for later, TODO revisit
}

int LocalSocket::availableBytesForReading()
{
    int available = 0;
    if (ioctl(m_fd, FIONREAD, &available) < 0) {
        available = 0;
    }
    return available;
}

array LocalSocket::read(byte *buffer, int maxSize)
{
    array ret;
    if (maxSize <= 0) {
        return ret;
    }

    // recvmsg-with-control-message boilerplate
    struct msghdr recv_msg;
    char cmsgBuf[CMSG_SPACE(sizeof(int) * maxFds)];
    memset(cmsgBuf, 0, sizeof(cmsgBuf));

    recv_msg.msg_control = cmsgBuf;
    recv_msg.msg_controllen = sizeof(cmsgBuf);
    recv_msg.msg_name = 0;
    recv_msg.msg_namelen = 0;
    recv_msg.msg_flags = 0;

    struct iovec iov;
    recv_msg.msg_iov = &iov;
    recv_msg.msg_iovlen = 1;

    // end boilerplate

    ret.begin = buffer;
    ret.length = 0;
    iov.iov_base = ret.begin;
    iov.iov_len = maxSize;
    while (iov.iov_len > 0) {
        int nbytes =  recvmsg(m_fd, &recv_msg, 0);
        if (nbytes < 0 && errno != EINTR) {
            close();
            return ret;
        }
        ret.length += nbytes;
        iov.iov_base = static_cast<char *>(iov.iov_base) + nbytes;
        iov.iov_len -= nbytes;
    }

    // done reading "regular data", now read any file descriptors passed via control messages

    struct cmsghdr *c_msg = CMSG_FIRSTHDR(&recv_msg);
    if (c_msg && c_msg->cmsg_level == SOL_SOCKET && c_msg->cmsg_type == SCM_RIGHTS) {
        const int len = c_msg->cmsg_len / sizeof(int);
        int *data = reinterpret_cast<int *>(CMSG_DATA(c_msg));
        for (int i = 0; i < len; i++) {
            // TODO
            // message.appendFileDescriptor(data[i]);
        }
    }

    return ret;
}

bool LocalSocket::isOpen()
{
    return m_fd != -1;
}

int LocalSocket::fileDescriptor() const
{
    return m_fd;
}

void LocalSocket::notifyRead()
{
    if (availableBytesForReading()) {
        IConnection::notifyRead();
    } else {
        // This should really only happen in error cases! ### TODO test?
        close();
    }
}
<commit_msg>Non-blocking reading and writing in LocalSocket.<commit_after>#include "localsocket.h"

#include <errno.h>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include "sys/uio.h"
#include <sys/un.h>
#include <unistd.h>

#include <cassert>
#include <climits>
#include <cstdlib>
#include <cstring>

// HACK, put this somewhere else (get the value from original d-bus? or is it infinite?)
static const int maxFds = 12;

using namespace std;

LocalSocket::LocalSocket(int fd)
   : m_fd(fd)
{}

LocalSocket::LocalSocket(const string &socketFilePath)
   : m_fd(-1)
{
    const int fd = socket(PF_UNIX, SOCK_STREAM, 0);
    if (fd < 0) {
        return;
    }
    // don't let forks inherit the file descriptor - that can cause confusion...
    fcntl(fd, F_SETFD, FD_CLOEXEC);

    struct sockaddr_un addr;
    addr.sun_family = PF_UNIX;
    bool ok = socketFilePath.length() < sizeof(addr.sun_path);
    if (ok) {
        memcpy(addr.sun_path, socketFilePath.c_str(), socketFilePath.length());
    }

    ok = ok && (connect(fd, (struct sockaddr *)&addr, sizeof(sa_family_t) + socketFilePath.length()) == 0);

    if (ok) {
        m_fd = fd;
    } else {
        ::close(fd);
    }
}

LocalSocket::~LocalSocket()
{
    close();
}

void LocalSocket::close()
{
    setEventDispatcher(0);
    if (m_fd >= 0) {
        ::close(m_fd);
    }
    m_fd = -1;
}

int LocalSocket::write(array a)
{
    if (m_fd < 0) {
        return 0; // TODO -1?
    }

    // sendmsg  boilerplate
    struct msghdr send_msg;
    struct iovec iov;

    send_msg.msg_name = 0;
    send_msg.msg_namelen = 0;
    send_msg.msg_flags = 0;
    send_msg.msg_iov = &iov;
    send_msg.msg_iovlen = 1;

    iov.iov_base = a.begin;
    iov.iov_len = a.length;

    // we can only send a fixed number of fds anyway due to the non-flexible size of the control message
    // receive buffer, so we set an arbitrary limit.
    const int numFds = 0; // TODO - how many should we get? do we need to know?
    assert(numFds <= maxFds);

    char cmsgBuf[CMSG_SPACE(sizeof(int) * maxFds)];

    if (numFds) {
        // fill in a control message
        send_msg.msg_control = cmsgBuf;
        send_msg.msg_controllen = CMSG_SPACE(sizeof(int) * numFds);

        struct cmsghdr *c_msg = CMSG_FIRSTHDR(&send_msg);
        c_msg->cmsg_len = CMSG_LEN(sizeof(int) * numFds);
        c_msg->cmsg_level = SOL_SOCKET;
        c_msg->cmsg_type = SCM_RIGHTS;

        // set the control data to pass - this is why we don't use the simpler write()
        for (int i = 0; i < numFds; i++) {
            // TODO
            // reinterpret_cast<int *>(CMSG_DATA(c_msg))[i] = message.m_fileDescriptors[i];
        }
    } else {
        // no file descriptor to send, no control message
        send_msg.msg_control = 0;
        send_msg.msg_controllen = 0;
    }

    while (iov.iov_len > 0) {
        int nbytes = sendmsg(m_fd, &send_msg, MSG_DONTWAIT);
        if (nbytes < 0) {
            if (errno == EINTR) {
                continue;
            }
            // if we were notified for writing, we must have written at least one byte before getting
            // EAGAIN aka EWOULDBLOCK
            if (errno == EAGAIN && iov.iov_len < a.length) {
                break;
            }
            close();
            return false;
        }

        iov.iov_base = static_cast<char *>(iov.iov_base) + nbytes;
        iov.iov_len -= nbytes;
        // sendmsg() should always send the number of bytes asked for or block, so...
        if (nbytes != 0) {
            assert(iov.iov_len == 0);
        }
    }

    return a.length - iov.iov_len; // "- iov.iov_len" is for later, TODO revisit
}

int LocalSocket::availableBytesForReading()
{
    int available = 0;
    if (ioctl(m_fd, FIONREAD, &available) < 0) {
        available = 0;
    }
    return available;
}

array LocalSocket::read(byte *buffer, int maxSize)
{
    array ret;
    if (maxSize <= 0) {
        return ret;
    }

    // recvmsg-with-control-message boilerplate
    struct msghdr recv_msg;
    char cmsgBuf[CMSG_SPACE(sizeof(int) * maxFds)];
    memset(cmsgBuf, 0, sizeof(cmsgBuf));

    recv_msg.msg_control = cmsgBuf;
    recv_msg.msg_controllen = sizeof(cmsgBuf);
    recv_msg.msg_name = 0;
    recv_msg.msg_namelen = 0;
    recv_msg.msg_flags = 0;

    struct iovec iov;
    recv_msg.msg_iov = &iov;
    recv_msg.msg_iovlen = 1;

    // end boilerplate

    ret.begin = buffer;
    ret.length = 0;
    iov.iov_base = ret.begin;
    iov.iov_len = maxSize;
    while (iov.iov_len > 0) {
        int nbytes = recvmsg(m_fd, &recv_msg, MSG_DONTWAIT);
        if (nbytes < 0) {
            if (errno == EINTR) {
                continue;
            }
            // if we were notified for reading, we must have read at least one byte before getting
            // EAGAIN aka EWOULDBLOCK
            if (errno == EAGAIN && iov.iov_len < maxSize) {
                break;
            }
            close();
            return ret;
        }
        ret.length += nbytes;
        iov.iov_base = static_cast<char *>(iov.iov_base) + nbytes;
        iov.iov_len -= nbytes;
    }

    // done reading "regular data", now read any file descriptors passed via control messages

    struct cmsghdr *c_msg = CMSG_FIRSTHDR(&recv_msg);
    if (c_msg && c_msg->cmsg_level == SOL_SOCKET && c_msg->cmsg_type == SCM_RIGHTS) {
        const int len = c_msg->cmsg_len / sizeof(int);
        int *data = reinterpret_cast<int *>(CMSG_DATA(c_msg));
        for (int i = 0; i < len; i++) {
            // TODO
            // message.appendFileDescriptor(data[i]);
        }
    }

    return ret;
}

bool LocalSocket::isOpen()
{
    return m_fd != -1;
}

int LocalSocket::fileDescriptor() const
{
    return m_fd;
}

void LocalSocket::notifyRead()
{
    if (availableBytesForReading()) {
        IConnection::notifyRead();
    } else {
        // This should really only happen in error cases! ### TODO test?
        close();
    }
}
<|endoftext|>