<commit_before>/*
 * Copyright (c) 2016 Fastly, Inc.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to
 * deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 * sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 * IN THE SOFTWARE.
 */

/*
 * This file implements a test harness for using h2o with LibFuzzer.
 * See http://llvm.org/docs/LibFuzzer.html for more info.
 */

#define H2O_USE_EPOLL 1
#include <string.h>
#include <errno.h>
#include <limits.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/select.h>
#include <sys/wait.h>
#include <unistd.h>
#include <fcntl.h>

#include "h2o.h"
#include "h2o/http1.h"
#include "h2o/http2.h"
#include "h2o/url.h"
#include "h2o/memcached.h"

#if !defined(HTTP1) && !defined(HTTP2)
#error "Please defined one of HTTP1 or HTTP2"
#endif

#if defined(HTTP1) && defined(HTTP2)
#error "Please defined one of HTTP1 or HTTP2, but not both"
#endif

static h2o_globalconf_t config;
static h2o_context_t ctx;
static h2o_accept_ctx_t accept_ctx;
static int client_timeout_ms;
static char unix_listener[PATH_MAX];

/*
 * Registers a request handler with h2o
 */
static h2o_pathconf_t *register_handler(h2o_hostconf_t *hostconf, const char *path, int (*on_req)(h2o_handler_t *, h2o_req_t *))
{
    h2o_pathconf_t *pathconf = h2o_config_register_path(hostconf, path, 0);
    h2o_handler_t *handler = h2o_create_handler(pathconf, sizeof(*handler));
    handler->on_req = on_req;
    return pathconf;
}

/*
 * Request handler used for testing. Returns a basic "200 OK" response.
 */
static int chunked_test(h2o_handler_t *self, h2o_req_t *req)
{
    static h2o_generator_t generator = {NULL, NULL};

    if (!h2o_memis(req->method.base, req->method.len, H2O_STRLIT("GET")))
        return -1;

    h2o_iovec_t body = h2o_strdup(&req->pool, "hello world\n", SIZE_MAX);
    req->res.status = 200;
    req->res.reason = "OK";
    h2o_add_header(&req->pool, &req->res.headers, H2O_TOKEN_CONTENT_TYPE, NULL, H2O_STRLIT("text/plain"));
    h2o_start_response(req, &generator);
    h2o_send(req, &body, 1, H2O_SEND_STATE_FINAL);

    return 0;
}

/* copy from src to dst, return true if src has EOF */
static int drain(int fd)
{
    char buf[4096];
    ssize_t n;

    n = read(fd, buf, sizeof(buf));
    if (n <= 0) {
        return 1;
    }
    return 0;
}

/* A request sent from client thread to h2o server */
struct writer_thread_arg {
    char *buf;
    size_t len;
    int fd;
    h2o_barrier_t barrier;
};

/*
 * Reads writer_thread_arg from fd and stores to buf
 */
static void read_fully(int fd, char *buf, size_t len)
{
    int done = 0;
    while (len) {
        int ret;
        while ((ret = read(fd, buf + done, len)) == -1 && errno == EINTR)
            ;
        if (ret <= 0) {
            abort();
        }
        done += ret;
        len -= ret;
    }
}

/*
 * Writes the writer_thread_args at buf to fd
 */
static void write_fully(int fd, char *buf, size_t len, int abort_on_err)
{
    int done = 0;
    while (len) {
        int ret;
        while ((ret = write(fd, buf + done, len)) == -1 && errno == EINTR)
            ;
        if (ret <= 0) {
            if (abort_on_err)
                abort();
            else
                return;
        }
        done += ret;
        len -= ret;
    }
}

#define OK_RESP                                                                                \
    "HTTP/1.0 200 OK\r\n"                                                                                \
    "Connection: Close\r\n\r\nOk"
#define OK_RESP_LEN (sizeof(OK_RESP) - 1)

void *upstream_thread(void *arg)
{
    char *dirname = (char *)arg;
    char path[PATH_MAX];
    char rbuf[1 * 1024 * 1024];
    snprintf(path, sizeof(path), "/%s/_.sock", dirname);
    int sd = socket(AF_UNIX, SOCK_STREAM, 0);
    if (sd < 0) {
        abort();
    }
    struct sockaddr_un addr;
    memset(&addr, 0, sizeof(addr));
    addr.sun_family = AF_UNIX;
    strncpy(addr.sun_path, path, sizeof(addr.sun_path) - 1);
    if (bind(sd, (struct sockaddr *)&addr, sizeof(addr)) != 0) {
        abort();
    }
    if (listen(sd, 100) != 0) {
        abort();
    }

    while (1) {
        struct sockaddr_un caddr;
        socklen_t slen = 0;
        int cfs = accept(sd, (struct sockaddr *)&caddr, &slen);
        if (cfs < 0) {
            continue;
        }
        read(cfs, rbuf, sizeof(rbuf));
        write_fully(cfs, (char *)OK_RESP, OK_RESP_LEN, 0);
        close(cfs);
    }
}
/*
 * Thread: Loops writing fuzzed req to socket and then reading results back.
 * Acts as a client to h2o. *arg points to file descripter to read
 * writer_thread_args from.
 */
void *writer_thread(void *arg)
{
    int rfd = (long)arg;
    while (1) {
        int pos, sockinp, sockoutp, cnt, len;
        char *buf;
        struct writer_thread_arg *wta;

        /* Get fuzzed request */
        read_fully(rfd, (char *)&wta, sizeof(wta));

        pos = 0;
        sockinp = wta->fd;
        sockoutp = wta->fd;
        cnt = 0;
        buf = wta->buf;
        len = wta->len;

        /*
         * Send fuzzed req and read results until the socket is closed (or
         * something spurious happens)
         */
        while (cnt++ < 20 && (pos < len || sockinp >= 0)) {
#define MARKER "\n--MARK--\n"
            /* send 1 packet */
            if (pos < len) {
                char *p = (char *)memmem(buf + pos, len - pos, MARKER, sizeof(MARKER) - 1);
                if (p) {
                    int l = p - (buf + pos);
                    write(sockoutp, buf + pos, l);
                    pos += l;
                    pos += sizeof(MARKER) - 1;
                }
            } else {
                if (sockinp >= 0) {
                    shutdown(sockinp, SHUT_WR);
                }
            }

            /* drain socket */
            if (sockinp >= 0) {
                struct timeval timeo;
                fd_set rd;
                int n;

                FD_ZERO(&rd);
                FD_SET(sockinp, &rd);
                timeo.tv_sec = 0;
                timeo.tv_usec = client_timeout_ms * 1000;
                n = select(sockinp + 1, &rd, NULL, NULL, &timeo);
                if (n > 0 && FD_ISSET(sockinp, &rd) && drain(sockinp)) {
                    sockinp = -1;
                }
            }
        }
        close(wta->fd);
        h2o_barrier_wait(&wta->barrier);
        h2o_barrier_destroy(&wta->barrier);
        free(wta);
    }
}

/*
 * Creates socket pair and passes fuzzed req to a thread (the HTTP[/2] client)
 * for writing to the target h2o server. Returns the server socket fd.
 */
static int feeder(int sfd, char *buf, size_t len, h2o_barrier_t **barrier)
{
    int pair[2];
    struct writer_thread_arg *wta;

    if (socketpair(AF_UNIX, SOCK_STREAM, 0, pair) == -1)
        return -1;

    wta = (struct writer_thread_arg *)malloc(sizeof(*wta));
    wta->fd = pair[0];
    wta->buf = buf;
    wta->len = len;
    h2o_barrier_init(&wta->barrier, 2);
    *barrier = &wta->barrier;

    write_fully(sfd, (char *)&wta, sizeof(wta), 1);
    return pair[1];
}

/*
 * Creates/connects socket pair for client/server interaction and passes
 * fuzzed request to client for sending.
 * Returns server socket fd.
 */
static int create_accepted(int sfd, char *buf, size_t len, h2o_barrier_t **barrier)
{
    int fd;
    h2o_socket_t *sock;
    struct timeval connected_at = h2o_gettimeofday(ctx.loop);

    /* Create an HTTP[/2] client that will send the fuzzed request */
    fd = feeder(sfd, buf, len, barrier);
    if (fd < 0) {
        abort();
    }

    /* Pass the server socket to h2o and invoke request processing */
    sock = h2o_evloop_socket_create(ctx.loop, fd, H2O_SOCKET_FLAG_IS_ACCEPTED_CONNECTION);

#if defined(HTTP1)
    h2o_http1_accept(&accept_ctx, sock, connected_at);
#else
    h2o_http2_accept(&accept_ctx, sock, connected_at);
#endif

    return fd;
}

/*
 * Returns true if fd if valid. Used to determine when connection is closed.
 */
static int is_valid_fd(int fd)
{
    return fcntl(fd, F_GETFD) != -1 || errno != EBADF;
}

/*
 * Entry point for libfuzzer.
 * See http://llvm.org/docs/LibFuzzer.html for more info
 */
static int init_done;
static int job_queue[2];
extern "C" int LLVMFuzzerTestOneInput(const uint8_t *Data, size_t Size)
{
    int c;
    h2o_loop_t *loop;
    h2o_hostconf_t *hostconf;
    pthread_t twriter;
    pthread_t tupstream;

    /*
     * Perform one-time initialization
     */
    if (!init_done) {
        const char *client_timeout_ms_str;
        static char tmpname[] = "/tmp/h2o-fuzz-XXXXXX";
        char *dirname;
        h2o_url_t upstream;
        signal(SIGPIPE, SIG_IGN);

        dirname = mkdtemp(tmpname);
        snprintf(unix_listener, sizeof(unix_listener), "http://[unix://%s/_.sock]/proxy", dirname);
        if ((client_timeout_ms_str = getenv("H2O_FUZZER_CLIENT_TIMEOUT")) != NULL)
            client_timeout_ms = atoi(client_timeout_ms_str);
        if (!client_timeout_ms)
            client_timeout_ms = 10;

        /* Create a single h2o host with multiple request handlers */
        h2o_config_init(&config);
        config.http2.idle_timeout = 10 * 1000;
        config.http1.req_timeout = 10 * 1000;
        config.proxy.io_timeout = 10 * 1000;
        h2o_proxy_config_vars_t proxy_config = {};
        proxy_config.io_timeout = 10 * 1000;
        hostconf = h2o_config_register_host(&config, h2o_iovec_init(H2O_STRLIT(unix_listener)), 65535);
        register_handler(hostconf, "/chunked-test", chunked_test);
        h2o_url_parse(unix_listener, strlen(unix_listener), &upstream);
        h2o_socketpool_t *sockpool = new h2o_socketpool_t();
        h2o_socketpool_target_t *target = h2o_socketpool_create_target(&upstream, NULL);
        h2o_socketpool_init_specific(sockpool, SIZE_MAX /* FIXME */, &target, 1, NULL);
        h2o_socketpool_set_timeout(sockpool, 2000);
        h2o_socketpool_set_ssl_ctx(sockpool, NULL);
        h2o_proxy_register_reverse_proxy(h2o_config_register_path(hostconf, "/reproxy-test", 0), &proxy_config, sockpool);
        h2o_file_register(h2o_config_register_path(hostconf, "/", 0), "./examples/doc_root", NULL, NULL, 0);

        loop = h2o_evloop_create();
        h2o_context_init(&ctx, loop, &config);

        accept_ctx.ctx = &ctx;
        accept_ctx.hosts = config.hosts;

        /* Create a thread to act as the HTTP client */
        if (socketpair(AF_UNIX, SOCK_STREAM, 0, job_queue) != 0) {
            abort();
        }
        if (pthread_create(&twriter, NULL, writer_thread, (void *)(long)job_queue[1]) != 0) {
            abort();
        }
        if (pthread_create(&tupstream, NULL, upstream_thread, dirname) != 0) {
            abort();
        }
        init_done = 1;
    }

    /*
     * Pass fuzzed request to client thread and get h2o server socket for
     * use below
     */
    h2o_barrier_t *end;
    c = create_accepted(job_queue[0], (char *)Data, (size_t)Size, &end);
    if (c < 0) {
        goto Error;
    }

    /* Loop until the connection is closed by the client or server */
    while (is_valid_fd(c)) {
        h2o_evloop_run(ctx.loop, 10);
    }

    h2o_barrier_wait(end);
    return 0;
Error:
    return 1;
}
<commit_msg>fuzzer: fix init time race<commit_after>/*
 * Copyright (c) 2016 Fastly, Inc.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to
 * deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 * sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 * IN THE SOFTWARE.
 */

/*
 * This file implements a test harness for using h2o with LibFuzzer.
 * See http://llvm.org/docs/LibFuzzer.html for more info.
 */

#define H2O_USE_EPOLL 1
#include <string.h>
#include <errno.h>
#include <limits.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/select.h>
#include <sys/wait.h>
#include <unistd.h>
#include <fcntl.h>

#include "h2o.h"
#include "h2o/http1.h"
#include "h2o/http2.h"
#include "h2o/url.h"
#include "h2o/memcached.h"

#if !defined(HTTP1) && !defined(HTTP2)
#error "Please defined one of HTTP1 or HTTP2"
#endif

#if defined(HTTP1) && defined(HTTP2)
#error "Please defined one of HTTP1 or HTTP2, but not both"
#endif

static h2o_globalconf_t config;
static h2o_context_t ctx;
static h2o_accept_ctx_t accept_ctx;
static int client_timeout_ms;
static char unix_listener[PATH_MAX];

/*
 * Registers a request handler with h2o
 */
static h2o_pathconf_t *register_handler(h2o_hostconf_t *hostconf, const char *path, int (*on_req)(h2o_handler_t *, h2o_req_t *))
{
    h2o_pathconf_t *pathconf = h2o_config_register_path(hostconf, path, 0);
    h2o_handler_t *handler = h2o_create_handler(pathconf, sizeof(*handler));
    handler->on_req = on_req;
    return pathconf;
}

/*
 * Request handler used for testing. Returns a basic "200 OK" response.
 */
static int chunked_test(h2o_handler_t *self, h2o_req_t *req)
{
    static h2o_generator_t generator = {NULL, NULL};

    if (!h2o_memis(req->method.base, req->method.len, H2O_STRLIT("GET")))
        return -1;

    h2o_iovec_t body = h2o_strdup(&req->pool, "hello world\n", SIZE_MAX);
    req->res.status = 200;
    req->res.reason = "OK";
    h2o_add_header(&req->pool, &req->res.headers, H2O_TOKEN_CONTENT_TYPE, NULL, H2O_STRLIT("text/plain"));
    h2o_start_response(req, &generator);
    h2o_send(req, &body, 1, H2O_SEND_STATE_FINAL);

    return 0;
}

/* copy from src to dst, return true if src has EOF */
static int drain(int fd)
{
    char buf[4096];
    ssize_t n;

    n = read(fd, buf, sizeof(buf));
    if (n <= 0) {
        return 1;
    }
    return 0;
}

/* A request sent from client thread to h2o server */
struct writer_thread_arg {
    char *buf;
    size_t len;
    int fd;
    h2o_barrier_t barrier;
};

/*
 * Reads writer_thread_arg from fd and stores to buf
 */
static void read_fully(int fd, char *buf, size_t len)
{
    int done = 0;
    while (len) {
        int ret;
        while ((ret = read(fd, buf + done, len)) == -1 && errno == EINTR)
            ;
        if (ret <= 0) {
            abort();
        }
        done += ret;
        len -= ret;
    }
}

/*
 * Writes the writer_thread_args at buf to fd
 */
static void write_fully(int fd, char *buf, size_t len, int abort_on_err)
{
    int done = 0;
    while (len) {
        int ret;
        while ((ret = write(fd, buf + done, len)) == -1 && errno == EINTR)
            ;
        if (ret <= 0) {
            if (abort_on_err)
                abort();
            else
                return;
        }
        done += ret;
        len -= ret;
    }
}

#define OK_RESP                                                                                \
    "HTTP/1.0 200 OK\r\n"                                                                                \
    "Connection: Close\r\n\r\nOk"
#define OK_RESP_LEN (sizeof(OK_RESP) - 1)

static h2o_barrier_t init_barrier;
void *upstream_thread(void *arg)
{
    char *dirname = (char *)arg;
    char path[PATH_MAX];
    char rbuf[1 * 1024 * 1024];
    snprintf(path, sizeof(path), "/%s/_.sock", dirname);
    int sd = socket(AF_UNIX, SOCK_STREAM, 0);
    if (sd < 0) {
        abort();
    }
    struct sockaddr_un addr;
    memset(&addr, 0, sizeof(addr));
    addr.sun_family = AF_UNIX;
    strncpy(addr.sun_path, path, sizeof(addr.sun_path) - 1);
    if (bind(sd, (struct sockaddr *)&addr, sizeof(addr)) != 0) {
        abort();
    }
    if (listen(sd, 100) != 0) {
        abort();
    }

    h2o_barrier_wait(&init_barrier);
    while (1) {
        struct sockaddr_un caddr;
        socklen_t slen = 0;
        int cfs = accept(sd, (struct sockaddr *)&caddr, &slen);
        if (cfs < 0) {
            continue;
        }
        read(cfs, rbuf, sizeof(rbuf));
        write_fully(cfs, (char *)OK_RESP, OK_RESP_LEN, 0);
        close(cfs);
    }
}
/*
 * Thread: Loops writing fuzzed req to socket and then reading results back.
 * Acts as a client to h2o. *arg points to file descripter to read
 * writer_thread_args from.
 */
void *writer_thread(void *arg)
{
    int rfd = (long)arg;
    while (1) {
        int pos, sockinp, sockoutp, cnt, len;
        char *buf;
        struct writer_thread_arg *wta;

        /* Get fuzzed request */
        read_fully(rfd, (char *)&wta, sizeof(wta));

        pos = 0;
        sockinp = wta->fd;
        sockoutp = wta->fd;
        cnt = 0;
        buf = wta->buf;
        len = wta->len;

        /*
         * Send fuzzed req and read results until the socket is closed (or
         * something spurious happens)
         */
        while (cnt++ < 20 && (pos < len || sockinp >= 0)) {
#define MARKER "\n--MARK--\n"
            /* send 1 packet */
            if (pos < len) {
                char *p = (char *)memmem(buf + pos, len - pos, MARKER, sizeof(MARKER) - 1);
                if (p) {
                    int l = p - (buf + pos);
                    write(sockoutp, buf + pos, l);
                    pos += l;
                    pos += sizeof(MARKER) - 1;
                }
            } else {
                if (sockinp >= 0) {
                    shutdown(sockinp, SHUT_WR);
                }
            }

            /* drain socket */
            if (sockinp >= 0) {
                struct timeval timeo;
                fd_set rd;
                int n;

                FD_ZERO(&rd);
                FD_SET(sockinp, &rd);
                timeo.tv_sec = 0;
                timeo.tv_usec = client_timeout_ms * 1000;
                n = select(sockinp + 1, &rd, NULL, NULL, &timeo);
                if (n > 0 && FD_ISSET(sockinp, &rd) && drain(sockinp)) {
                    sockinp = -1;
                }
            }
        }
        close(wta->fd);
        h2o_barrier_wait(&wta->barrier);
        h2o_barrier_destroy(&wta->barrier);
        free(wta);
    }
}

/*
 * Creates socket pair and passes fuzzed req to a thread (the HTTP[/2] client)
 * for writing to the target h2o server. Returns the server socket fd.
 */
static int feeder(int sfd, char *buf, size_t len, h2o_barrier_t **barrier)
{
    int pair[2];
    struct writer_thread_arg *wta;

    if (socketpair(AF_UNIX, SOCK_STREAM, 0, pair) == -1)
        return -1;

    wta = (struct writer_thread_arg *)malloc(sizeof(*wta));
    wta->fd = pair[0];
    wta->buf = buf;
    wta->len = len;
    h2o_barrier_init(&wta->barrier, 2);
    *barrier = &wta->barrier;

    write_fully(sfd, (char *)&wta, sizeof(wta), 1);
    return pair[1];
}

/*
 * Creates/connects socket pair for client/server interaction and passes
 * fuzzed request to client for sending.
 * Returns server socket fd.
 */
static int create_accepted(int sfd, char *buf, size_t len, h2o_barrier_t **barrier)
{
    int fd;
    h2o_socket_t *sock;
    struct timeval connected_at = h2o_gettimeofday(ctx.loop);

    /* Create an HTTP[/2] client that will send the fuzzed request */
    fd = feeder(sfd, buf, len, barrier);
    if (fd < 0) {
        abort();
    }

    /* Pass the server socket to h2o and invoke request processing */
    sock = h2o_evloop_socket_create(ctx.loop, fd, H2O_SOCKET_FLAG_IS_ACCEPTED_CONNECTION);

#if defined(HTTP1)
    h2o_http1_accept(&accept_ctx, sock, connected_at);
#else
    h2o_http2_accept(&accept_ctx, sock, connected_at);
#endif

    return fd;
}

/*
 * Returns true if fd if valid. Used to determine when connection is closed.
 */
static int is_valid_fd(int fd)
{
    return fcntl(fd, F_GETFD) != -1 || errno != EBADF;
}

/*
 * Entry point for libfuzzer.
 * See http://llvm.org/docs/LibFuzzer.html for more info
 */
static int init_done;
static int job_queue[2];
extern "C" int LLVMFuzzerTestOneInput(const uint8_t *Data, size_t Size)
{
    int c;
    h2o_loop_t *loop;
    h2o_hostconf_t *hostconf;
    pthread_t twriter;
    pthread_t tupstream;

    /*
     * Perform one-time initialization
     */
    if (!init_done) {
        const char *client_timeout_ms_str;
        static char tmpname[] = "/tmp/h2o-fuzz-XXXXXX";
        char *dirname;
        h2o_url_t upstream;
        h2o_barrier_init(&init_barrier, 2);
        signal(SIGPIPE, SIG_IGN);

        dirname = mkdtemp(tmpname);
        snprintf(unix_listener, sizeof(unix_listener), "http://[unix://%s/_.sock]/proxy", dirname);
        if ((client_timeout_ms_str = getenv("H2O_FUZZER_CLIENT_TIMEOUT")) != NULL)
            client_timeout_ms = atoi(client_timeout_ms_str);
        if (!client_timeout_ms)
            client_timeout_ms = 10;

        /* Create a single h2o host with multiple request handlers */
        h2o_config_init(&config);
        config.http2.idle_timeout = 10 * 1000;
        config.http1.req_timeout = 10 * 1000;
        config.proxy.io_timeout = 10 * 1000;
        h2o_proxy_config_vars_t proxy_config = {};
        proxy_config.io_timeout = 10 * 1000;
        hostconf = h2o_config_register_host(&config, h2o_iovec_init(H2O_STRLIT(unix_listener)), 65535);
        register_handler(hostconf, "/chunked-test", chunked_test);
        h2o_url_parse(unix_listener, strlen(unix_listener), &upstream);
        h2o_socketpool_t *sockpool = new h2o_socketpool_t();
        h2o_socketpool_target_t *target = h2o_socketpool_create_target(&upstream, NULL);
        h2o_socketpool_init_specific(sockpool, SIZE_MAX /* FIXME */, &target, 1, NULL);
        h2o_socketpool_set_timeout(sockpool, 2000);
        h2o_socketpool_set_ssl_ctx(sockpool, NULL);
        h2o_proxy_register_reverse_proxy(h2o_config_register_path(hostconf, "/reproxy-test", 0), &proxy_config, sockpool);
        h2o_file_register(h2o_config_register_path(hostconf, "/", 0), "./examples/doc_root", NULL, NULL, 0);

        loop = h2o_evloop_create();
        h2o_context_init(&ctx, loop, &config);

        accept_ctx.ctx = &ctx;
        accept_ctx.hosts = config.hosts;

        /* Create a thread to act as the HTTP client */
        if (socketpair(AF_UNIX, SOCK_STREAM, 0, job_queue) != 0) {
            abort();
        }
        if (pthread_create(&twriter, NULL, writer_thread, (void *)(long)job_queue[1]) != 0) {
            abort();
        }
        if (pthread_create(&tupstream, NULL, upstream_thread, dirname) != 0) {
            abort();
        }
        h2o_barrier_wait(&init_barrier);
        init_done = 1;
    }

    /*
     * Pass fuzzed request to client thread and get h2o server socket for
     * use below
     */
    h2o_barrier_t *end;
    c = create_accepted(job_queue[0], (char *)Data, (size_t)Size, &end);
    if (c < 0) {
        goto Error;
    }

    /* Loop until the connection is closed by the client or server */
    while (is_valid_fd(c)) {
        h2o_evloop_run(ctx.loop, 10);
    }

    h2o_barrier_wait(end);
    return 0;
Error:
    return 1;
}
<|endoftext|>


<commit_before>/*
 * nwconv.hpp
 *  a class to convert between std::string and std::wstring
 *
 *  written by janus_wel<janus.wel.3@gmail.com>
 *  This source code is in the public domain, and has NO WARRANTY.
 * */

#ifndef NWCONV_HPP
#define NWCONV_HPP

#include <locale>
#include <string>
#include <vector>
#include <stdexcept>
#include "wexcept.hpp"

namespace util {
    namespace string {
        template<typename Internal, typename External, typename State>
            class basic_nwconv {
                private:
                    // typedefs
                    typedef std::codecvt<Internal, External, State> codecvt_type;
                    typedef typename codecvt_type::state_type       state_type;
                    typedef typename codecvt_type::intern_type      intern_type;
                    typedef typename codecvt_type::extern_type      extern_type;
                    typedef std::basic_string<intern_type>          intern_string_type;
                    typedef std::basic_string<extern_type>          extern_string_type;

                    // cache a reference and objects
                    const codecvt_type& codecvt(void) const {
                        static const codecvt_type& conv = std::use_facet<codecvt_type>(loc);
                        return conv;
                    }
                    state_type& in_state(void) const {
                        static state_type in_stat;
                        return in_stat;
                    }
                    state_type& out_state(void) const {
                        static state_type out_stat;
                        return out_stat;
                    }

                    // member variables
                    std::locale loc;

                public:
                    // constructor
                    explicit basic_nwconv(std::locale loc = std::locale("")) : loc(loc) {
                        if (!std::has_facet<codecvt_type>(loc))
                            throw std::logic_error("Specified locale doesn't have codecvt facet.");
                    }

                    // narrow to wide
                    intern_string_type ntow(const extern_string_type& src) const {
                        // preparations
                        // a size of src string
                        const std::size_t size = src.length();
                        // buffers
                        std::vector<intern_type> dst_vctr(size);
                        intern_type* const dst = &dst_vctr[0];
                        // a dummy variable and a pointer to the next position
                        const extern_type* dummy;
                        intern_type* next;
                        // cache the first position of c style characters
                        const extern_type* const s = src.c_str();

                        // do it
                        if (codecvt().in(
                                    in_state(),
                                    s,   s   + size, dummy,
                                    dst, dst + size, next) == codecvt_type::ok) {
                            return std::wstring(dst, next - dst);
                        }

                        std::string errmsg("failed: ");
                        errmsg.append(src);
                        throw std::logic_error(errmsg);
                    }

                    // wide to narrow
                    extern_string_type wton(const intern_string_type& src) const {
                        // preparations
                        // a size of src string
                        const std::size_t src_size = src.length();
                        // calc a size of dst string
                        // This value is expected maximum.
                        const std::size_t dst_size = codecvt().max_length() * src_size;
                        // buffers
                        std::vector<extern_type> dst_vctr(dst_size);
                        extern_type* const dst = &dst_vctr[0];
                        // a dummy variable and a pointer to the next position
                        const intern_type* dummy;
                        extern_type* next;
                        // cache the first position of c style characters
                        const intern_type* const s = src.c_str();

                        // do it
                        if (codecvt().out(
                                    out_state(),
                                    s,   s   + src_size, dummy,
                                    dst, dst + dst_size, next) == codecvt_type::ok) {
                            return std::string(dst, next - dst);
                        }

                        std::wstring errmsg(L"failed: ");
                        errmsg.append(src);
                        throw util::exception::wlogic_error(errmsg);
                    }
            };

        // for convenience
        typedef basic_nwconv<wchar_t, char, mbstate_t> nwconv;
    }
}

#endif // NWCONV_HPP

<commit_msg>Fix crashes when passing null strings<commit_after>/*
 * nwconv.hpp
 *  a class to convert between std::string and std::wstring
 *
 *  written by janus_wel<janus.wel.3@gmail.com>
 *  This source code is in the public domain, and has NO WARRANTY.
 * */

#ifndef NWCONV_HPP
#define NWCONV_HPP

#include <locale>
#include <string>
#include <vector>
#include <stdexcept>
#include "wexcept.hpp"

namespace util {
    namespace string {
        template<typename Internal, typename External, typename State>
            class basic_nwconv {
                private:
                    // typedefs
                    typedef std::codecvt<Internal, External, State> codecvt_type;
                    typedef typename codecvt_type::state_type       state_type;
                    typedef typename codecvt_type::intern_type      intern_type;
                    typedef typename codecvt_type::extern_type      extern_type;
                    typedef std::basic_string<intern_type>          intern_string_type;
                    typedef std::basic_string<extern_type>          extern_string_type;

                    // cache a reference and objects
                    const codecvt_type& codecvt(void) const {
                        static const codecvt_type& conv = std::use_facet<codecvt_type>(loc);
                        return conv;
                    }
                    state_type& in_state(void) const {
                        static state_type in_stat;
                        return in_stat;
                    }
                    state_type& out_state(void) const {
                        static state_type out_stat;
                        return out_stat;
                    }

                    // member variables
                    std::locale loc;

                public:
                    // constructor
                    explicit basic_nwconv(std::locale loc = std::locale("")) : loc(loc) {
                        if (!std::has_facet<codecvt_type>(loc))
                            throw std::logic_error("Specified locale doesn't have codecvt facet.");
                    }

                    // narrow to wide
                    intern_string_type ntow(const extern_string_type& src) const {
                        // preparations
                        // a size of src string
                        const std::size_t size = src.length();

                        // no more task
                        if (size == 0) return intern_string_type();

                        // buffers
                        std::vector<intern_type> dst_vctr(size);
                        intern_type* const dst = &dst_vctr[0];
                        // a dummy variable and a pointer to the next position
                        const extern_type* dummy;
                        intern_type* next;
                        // cache the first position of c style characters
                        const extern_type* const s = src.c_str();

                        // do it
                        if (codecvt().in(
                                    in_state(),
                                    s,   s   + size, dummy,
                                    dst, dst + size, next) == codecvt_type::ok) {
                            return std::wstring(dst, next - dst);
                        }

                        std::string errmsg("failed: ");
                        errmsg.append(src);
                        throw std::logic_error(errmsg);
                    }

                    // wide to narrow
                    extern_string_type wton(const intern_string_type& src) const {
                        // preparations
                        // a size of src string
                        const std::size_t src_size = src.length();

                        // no more task
                        if (src_size == 0) return extern_string_type();

                        // calc a size of dst string
                        // This value is expected maximum.
                        const std::size_t dst_size = codecvt().max_length() * src_size;
                        // buffers
                        std::vector<extern_type> dst_vctr(dst_size);
                        extern_type* const dst = &dst_vctr[0];
                        // a dummy variable and a pointer to the next position
                        const intern_type* dummy;
                        extern_type* next;
                        // cache the first position of c style characters
                        const intern_type* const s = src.c_str();

                        // do it
                        if (codecvt().out(
                                    out_state(),
                                    s,   s   + src_size, dummy,
                                    dst, dst + dst_size, next) == codecvt_type::ok) {
                            return std::string(dst, next - dst);
                        }

                        std::wstring errmsg(L"failed: ");
                        errmsg.append(src);
                        throw util::exception::wlogic_error(errmsg);
                    }
            };

        // for convenience
        typedef basic_nwconv<wchar_t, char, mbstate_t> nwconv;
    }
}

#endif // NWCONV_HPP

<|endoftext|>


<commit_before>// Copyright (c) 2006-2008 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "base/thread_local_storage.h"

#include <windows.h>

#include "base/logging.h"

// In order to make TLS destructors work, we need to keep function
// pointers to the destructor for each TLS that we allocate.
// We make this work by allocating a single OS-level TLS, which
// contains an array of slots for the application to use.  In
// parallel, we also allocate an array of destructors, which we
// keep track of and call when threads terminate.

// tls_key_ is the one native TLS that we use.  It stores our
// table.
long ThreadLocalStorage::tls_key_ = TLS_OUT_OF_INDEXES;

// tls_max_ is the high-water-mark of allocated thread local storage.
// We intentionally skip 0 so that it is not confused with an
// unallocated TLS slot.
long ThreadLocalStorage::tls_max_ = 1;

// An array of destructor function pointers for the slots.  If
// a slot has a destructor, it will be stored in its corresponding
// entry in this array.
ThreadLocalStorage::TLSDestructorFunc
  ThreadLocalStorage::tls_destructors_[kThreadLocalStorageSize];

void** ThreadLocalStorage::Initialize() {
  if (tls_key_ == TLS_OUT_OF_INDEXES) {
    long value = TlsAlloc();
    DCHECK(value != TLS_OUT_OF_INDEXES);

    // Atomically test-and-set the tls_key.  If the key is TLS_OUT_OF_INDEXES,
    // go ahead and set it.  Otherwise, do nothing, as another
    // thread already did our dirty work.
    if (InterlockedCompareExchange(&tls_key_, value, TLS_OUT_OF_INDEXES) !=
            TLS_OUT_OF_INDEXES) {
      // We've been shortcut. Another thread replaced tls_key_ first so we need
      // to destroy our index and use the one the other thread got first.
      TlsFree(value);
    }
  }
  DCHECK(TlsGetValue(tls_key_) == NULL);

  // Create an array to store our data.
  void** tls_data = new void*[kThreadLocalStorageSize];
  memset(tls_data, 0, sizeof(void*[kThreadLocalStorageSize]));
  TlsSetValue(tls_key_, tls_data);
  return tls_data;
}

ThreadLocalStorage::Slot::Slot(TLSDestructorFunc destructor)
    : initialized_(false) {
  Initialize(destructor);
}

bool ThreadLocalStorage::Slot::Initialize(TLSDestructorFunc destructor) {
  if (tls_key_ == TLS_OUT_OF_INDEXES || !TlsGetValue(tls_key_))
    ThreadLocalStorage::Initialize();

  // Grab a new slot.
  slot_ = InterlockedIncrement(&tls_max_) - 1;
  if (slot_ >= kThreadLocalStorageSize) {
    NOTREACHED();
    return false;
  }

  // Setup our destructor.
  tls_destructors_[slot_] = destructor;
  initialized_ = true;
  return true;
}

void ThreadLocalStorage::Slot::Free() {
  // At this time, we don't reclaim old indices for TLS slots.
  // So all we need to do is wipe the destructor.
  tls_destructors_[slot_] = NULL;
  initialized_ = false;
}

void* ThreadLocalStorage::Slot::Get() const {
  void** tls_data = static_cast<void**>(TlsGetValue(tls_key_));
  if (!tls_data)
    tls_data = ThreadLocalStorage::Initialize();
  DCHECK(slot_ >= 0 && slot_ < kThreadLocalStorageSize);
  return tls_data[slot_];
}

void ThreadLocalStorage::Slot::Set(void* value) {
  void** tls_data = static_cast<void**>(TlsGetValue(tls_key_));
  if (!tls_data)
    tls_data = ThreadLocalStorage::Initialize();
  DCHECK(slot_ >= 0 && slot_ < kThreadLocalStorageSize);
  tls_data[slot_] = value;
}

void ThreadLocalStorage::ThreadExit() {
  void** tls_data = static_cast<void**>(TlsGetValue(tls_key_));

  // Maybe we have never initialized TLS for this thread.
  if (!tls_data)
    return;

  for (int slot = 0; slot < tls_max_; slot++) {
    if (tls_destructors_[slot] != NULL) {
      void* value = tls_data[slot];
      tls_destructors_[slot](value);
    }
  }

  delete[] tls_data;

  // In case there are other "onexit" handlers...
  TlsSetValue(tls_key_, NULL);
}

// Thread Termination Callbacks.
// Windows doesn't support a per-thread destructor with its
// TLS primitives.  So, we build it manually by inserting a
// function to be called on each thread's exit.
// This magic is from http://www.codeproject.com/threads/tls.asp
// and it works for VC++ 7.0 and later.

#ifdef _WIN64

// This makes the linker create the TLS directory if it's not already
// there.  (e.g. if __declspec(thread) is not used).
#pragma comment(linker, "/INCLUDE:_tls_used")

#else  // _WIN64

// This makes the linker create the TLS directory if it's not already
// there.  (e.g. if __declspec(thread) is not used).
#pragma comment(linker, "/INCLUDE:__tls_used")

#endif  // _WIN64

// Static callback function to call with each thread termination.
void NTAPI OnThreadExit(PVOID module, DWORD reason, PVOID reserved)
{
  // On XP SP0 & SP1, the DLL_PROCESS_ATTACH is never seen. It is sent on SP2+
  // and on W2K and W2K3. So don't assume it is sent.
  if (DLL_THREAD_DETACH == reason || DLL_PROCESS_DETACH == reason)
    ThreadLocalStorage::ThreadExit();
}

// .CRT$XLA to .CRT$XLZ is an array of PIMAGE_TLS_CALLBACK pointers that are
// called automatically by the OS loader code (not the CRT) when the module is
// loaded and on thread creation. They are NOT called if the module has been
// loaded by a LoadLibrary() call. It must have implicitly been loaded at
// process startup.
// By implicitly loaded, I mean that it is directly referenced by the main EXE
// or by one of its dependent DLLs. Delay-loaded DLL doesn't count as being
// implicitly loaded.
//
// See VC\crt\src\tlssup.c for reference.
#ifdef _WIN64

// .CRT section is merged with .rdata on x64 so it must be constant data.
#pragma const_seg(".CRT$XLB")
// When defining a const variable, it must have external linkage to be sure the
// linker doesn't discard it. If this value is discarded, the OnThreadExit
// function will never be called.
extern const PIMAGE_TLS_CALLBACK p_thread_callback;
const PIMAGE_TLS_CALLBACK p_thread_callback = OnThreadExit;

// Reset the default section.
#pragma const_seg()

#else  // _WIN64

#pragma data_seg(".CRT$XLB")
PIMAGE_TLS_CALLBACK p_thread_callback = OnThreadExit;

// Reset the default section.
#pragma data_seg()

#endif  // _WIN64

<commit_msg>Don't get the TLS for invalid TLS indexes. This is making App Verifier complain. This case happens early when the process starts.<commit_after>// Copyright (c) 2006-2008 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "base/thread_local_storage.h"

#include <windows.h>

#include "base/logging.h"

// In order to make TLS destructors work, we need to keep function
// pointers to the destructor for each TLS that we allocate.
// We make this work by allocating a single OS-level TLS, which
// contains an array of slots for the application to use.  In
// parallel, we also allocate an array of destructors, which we
// keep track of and call when threads terminate.

// tls_key_ is the one native TLS that we use.  It stores our
// table.
long ThreadLocalStorage::tls_key_ = TLS_OUT_OF_INDEXES;

// tls_max_ is the high-water-mark of allocated thread local storage.
// We intentionally skip 0 so that it is not confused with an
// unallocated TLS slot.
long ThreadLocalStorage::tls_max_ = 1;

// An array of destructor function pointers for the slots.  If
// a slot has a destructor, it will be stored in its corresponding
// entry in this array.
ThreadLocalStorage::TLSDestructorFunc
  ThreadLocalStorage::tls_destructors_[kThreadLocalStorageSize];

void** ThreadLocalStorage::Initialize() {
  if (tls_key_ == TLS_OUT_OF_INDEXES) {
    long value = TlsAlloc();
    DCHECK(value != TLS_OUT_OF_INDEXES);

    // Atomically test-and-set the tls_key.  If the key is TLS_OUT_OF_INDEXES,
    // go ahead and set it.  Otherwise, do nothing, as another
    // thread already did our dirty work.
    if (InterlockedCompareExchange(&tls_key_, value, TLS_OUT_OF_INDEXES) !=
            TLS_OUT_OF_INDEXES) {
      // We've been shortcut. Another thread replaced tls_key_ first so we need
      // to destroy our index and use the one the other thread got first.
      TlsFree(value);
    }
  }
  DCHECK(TlsGetValue(tls_key_) == NULL);

  // Create an array to store our data.
  void** tls_data = new void*[kThreadLocalStorageSize];
  memset(tls_data, 0, sizeof(void*[kThreadLocalStorageSize]));
  TlsSetValue(tls_key_, tls_data);
  return tls_data;
}

ThreadLocalStorage::Slot::Slot(TLSDestructorFunc destructor)
    : initialized_(false) {
  Initialize(destructor);
}

bool ThreadLocalStorage::Slot::Initialize(TLSDestructorFunc destructor) {
  if (tls_key_ == TLS_OUT_OF_INDEXES || !TlsGetValue(tls_key_))
    ThreadLocalStorage::Initialize();

  // Grab a new slot.
  slot_ = InterlockedIncrement(&tls_max_) - 1;
  if (slot_ >= kThreadLocalStorageSize) {
    NOTREACHED();
    return false;
  }

  // Setup our destructor.
  tls_destructors_[slot_] = destructor;
  initialized_ = true;
  return true;
}

void ThreadLocalStorage::Slot::Free() {
  // At this time, we don't reclaim old indices for TLS slots.
  // So all we need to do is wipe the destructor.
  tls_destructors_[slot_] = NULL;
  initialized_ = false;
}

void* ThreadLocalStorage::Slot::Get() const {
  void** tls_data = static_cast<void**>(TlsGetValue(tls_key_));
  if (!tls_data)
    tls_data = ThreadLocalStorage::Initialize();
  DCHECK(slot_ >= 0 && slot_ < kThreadLocalStorageSize);
  return tls_data[slot_];
}

void ThreadLocalStorage::Slot::Set(void* value) {
  void** tls_data = static_cast<void**>(TlsGetValue(tls_key_));
  if (!tls_data)
    tls_data = ThreadLocalStorage::Initialize();
  DCHECK(slot_ >= 0 && slot_ < kThreadLocalStorageSize);
  tls_data[slot_] = value;
}

void ThreadLocalStorage::ThreadExit() {
  if (tls_key_ == TLS_OUT_OF_INDEXES)
    return;

  void** tls_data = static_cast<void**>(TlsGetValue(tls_key_));

  // Maybe we have never initialized TLS for this thread.
  if (!tls_data)
    return;

  for (int slot = 0; slot < tls_max_; slot++) {
    if (tls_destructors_[slot] != NULL) {
      void* value = tls_data[slot];
      tls_destructors_[slot](value);
    }
  }

  delete[] tls_data;

  // In case there are other "onexit" handlers...
  TlsSetValue(tls_key_, NULL);
}

// Thread Termination Callbacks.
// Windows doesn't support a per-thread destructor with its
// TLS primitives.  So, we build it manually by inserting a
// function to be called on each thread's exit.
// This magic is from http://www.codeproject.com/threads/tls.asp
// and it works for VC++ 7.0 and later.

#ifdef _WIN64

// This makes the linker create the TLS directory if it's not already
// there.  (e.g. if __declspec(thread) is not used).
#pragma comment(linker, "/INCLUDE:_tls_used")

#else  // _WIN64

// This makes the linker create the TLS directory if it's not already
// there.  (e.g. if __declspec(thread) is not used).
#pragma comment(linker, "/INCLUDE:__tls_used")

#endif  // _WIN64

// Static callback function to call with each thread termination.
void NTAPI OnThreadExit(PVOID module, DWORD reason, PVOID reserved)
{
  // On XP SP0 & SP1, the DLL_PROCESS_ATTACH is never seen. It is sent on SP2+
  // and on W2K and W2K3. So don't assume it is sent.
  if (DLL_THREAD_DETACH == reason || DLL_PROCESS_DETACH == reason)
    ThreadLocalStorage::ThreadExit();
}

// .CRT$XLA to .CRT$XLZ is an array of PIMAGE_TLS_CALLBACK pointers that are
// called automatically by the OS loader code (not the CRT) when the module is
// loaded and on thread creation. They are NOT called if the module has been
// loaded by a LoadLibrary() call. It must have implicitly been loaded at
// process startup.
// By implicitly loaded, I mean that it is directly referenced by the main EXE
// or by one of its dependent DLLs. Delay-loaded DLL doesn't count as being
// implicitly loaded.
//
// See VC\crt\src\tlssup.c for reference.
#ifdef _WIN64

// .CRT section is merged with .rdata on x64 so it must be constant data.
#pragma const_seg(".CRT$XLB")
// When defining a const variable, it must have external linkage to be sure the
// linker doesn't discard it. If this value is discarded, the OnThreadExit
// function will never be called.
extern const PIMAGE_TLS_CALLBACK p_thread_callback;
const PIMAGE_TLS_CALLBACK p_thread_callback = OnThreadExit;

// Reset the default section.
#pragma const_seg()

#else  // _WIN64

#pragma data_seg(".CRT$XLB")
PIMAGE_TLS_CALLBACK p_thread_callback = OnThreadExit;

// Reset the default section.
#pragma data_seg()

#endif  // _WIN64

<|endoftext|>


<commit_before>/*
 * Copyright (c) 2015 The Jackson Laboratory
 *
 * This software was developed by Gary Churchill's Lab at The Jackson
 * Laboratory (see http://research.jax.org/faculty/churchill).
 *
 * This is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This software 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 General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this software. If not, see <http://www.gnu.org/licenses/>.
 */

//
//  emase.cpp
//
//
//  Created by Glen Beane on 8/20/14.
//
//

#include <iostream>
#include <iomanip>
#include <cmath>
#include <getopt.h>

#include "alignment_incidence_matrix.h"
#include "sample_allelic_expression.h"
#include "python_interface.h"
#include "kallisto_import.h"

#define VERSION "0.2.4"

void print_help();

int main(int argc, char **argv)
{
    AlignmentIncidenceMatrix *aim;

    bool binary_input = false;
    int verbose = 0;

    int num_iterations;
    int max_iterations = 999;

    SampleAllelicExpression::model model = SampleAllelicExpression::MODEL_1;
    int m;

    clock_t t1, t2;
    float diff;

    std::string input_filename;
    std::string output_filename;
    std::string output_filename_counts;
    std::string transcript_length_file;
    std::string extension = ".pcl.bz2";
    std::string gene_file;

    //getopt related variables
    int c;
    int option_index = 0;
    bool bad_args = false;

    double tolerance = 0.0001;

    static struct option long_options[] = {
        {"help", no_argument, 0, 'h'},
        {"model", required_argument, 0, 'm'},
        {"outbase", required_argument, 0, 'o'},
        {"transcript-lengths", required_argument, 0, 'l'},
        {"max-iterations", required_argument, 0, 'i'},
        {"gene-mappings", required_argument, 0, 'g'},
        {"verbose", no_argument, &verbose, 1},
        {"version", no_argument, 0, 'V'},
        {"tolerance", required_argument, 0, 't'},
        {0, 0, 0, 0}
    };

    while ((c = getopt_long(argc, argv, "hm:o:l:i:bg:vVc:t:", long_options,
                            &option_index)) != -1) {
        switch (c) {
            case 'h':
                print_help();
                return 0;

            case 'm':
                m = std::stoi(optarg);
                if (m < 1 || m > 4) {
                    std::cerr << "[ERROR] Invalid model number specified. Valid options: 1, 2, 3, 4\n";
                }

                switch (m) {
                    case 1:
                        model = SampleAllelicExpression::MODEL_1;
                        break;
                    case 2:
                        model = SampleAllelicExpression::MODEL_2;
                        break;
                    case 3:
                        model = SampleAllelicExpression::MODEL_3;
                        break;
                    case 4:
                        model = SampleAllelicExpression::MODEL_4;
                        break;
                }

                break;

            case 'o':
                output_filename = std::string(optarg);
                output_filename_counts = output_filename.copy();
                output_filename.append(".tpm");
                output_filename_counts.append(".counts");
                break;

            case 'l':
                transcript_length_file = std::string(optarg);
                break;

            case 'i':
                max_iterations = std::stoi(optarg);
                break;

            case 'g':
                gene_file = std::string(optarg);
                break;

            case 'v':
                verbose = 1;
                break;

            case 'V':
               std::cout << VERSION << std::endl;
               return 0;

            case 't':
                tolerance = std::stod(optarg);
                break;

            case '?':
                bad_args = true;

        }
    }

    if (bad_args) {
        print_help();
        return 1;
    }

    if (argc - optind == 1) {
        input_filename = argv[optind];
    } else {
        std::cerr << "\n[ERROR] Missing required argument (input file name)\n";
        print_help();
        return 1;
    }

    if (output_filename.empty()) {
        //use default, based on the input file name but placed in current working directdory
        output_filename = "emase-zero.quantified.tpm";
        output_filename_counts = "emase-zero.quantified.counts";
    }

    std::cout << "\nemase-zero Version " << VERSION << std::endl <<std::endl;

    std::cout << "Alignment File: " << input_filename << std::endl;

    if (!gene_file.empty()) {
        std::cout << "Grouping File: " << gene_file << std::endl;
    }
    else {
        std::cout << "Grouping File: None\n";
    }

    if (!transcript_length_file.empty()) {
        std::cout << "Transcript Length File: " << transcript_length_file
                  << std::endl;
    }
    else {
        std::cout << "Transcript Length File: None\n";
    }

    std::cout << "EM Model: " << model << std::endl
              << "Output File (TPM): " << output_filename << std::endl
              << "Output File (Counts): " << output_filename_counts << std::endl
              << "----------------------------------------------------\n\n\n";

    std::cout << "Loading " << input_filename << "..." << std::endl;
    aim = loadFromBin(input_filename);

    if (!aim) {
        std::cerr << "Error loading binary input file\n";
        return 1;
    }

    std::vector<std::string> hap_names = aim->get_haplotype_names();

    if (verbose) {
        std::cout << "File had the following haplotype names:\n";
        for (auto it = hap_names.begin(); it != hap_names.end(); ++it) {
            std::cout << *it << "\t";
        }
        std::cout << std::endl;

        if (aim->has_equivalence_classes()) {
            std::cout << aim->num_alignment_classes()
                      << " alignment classes loaded ("
                      << aim->total_reads() << " total reads)\n";
        }
        else {
            std::cout << aim->total_reads() << " reads loaded\n";
        }
        std::cout << aim->num_transcripts() << " transcripts\n";

        std::cout << std::endl;
    }

    if (!transcript_length_file.empty()) {
        if (verbose) {
            std::cout << "Loading Transcript Length File "
                      << transcript_length_file << std::endl;
        }
        aim->loadTranscriptLengths(transcript_length_file);
    }

    if (!gene_file.empty()) {
        if (verbose) {
            std::cout << "Loading Gene Mapping File " << gene_file << std::endl;
        }
        aim->loadGeneMappings(gene_file);
    }

    if (model != SampleAllelicExpression::MODEL_4 && !aim->has_gene_mappings()) {
        if (!binary_input) {
            std::cerr << "[ERROR] File does not contain transcript to gene mapping information.  Only normalization Model 4 can be used.\n";
        }
        else {
            std::cerr << "[ERROR] Only model 4 is possible without gene to transcript information (--gene-mappings, -g). You specified another model.\n";
        }
        return 1;
    }

    t1 = clock();
    SampleAllelicExpression sae(aim, tolerance);
    t2 = clock();

    if (verbose) {
        diff = ((float)t2-(float)t1)/CLOCKS_PER_SEC;
        std::cout << "Time for initializing stack sum = " << diff << "s"
                  << std::endl;
    }

    if (max_iterations > 0) {
        num_iterations = 0;
        std::cout << "Beginning EM Iterations..." << std::endl;
        double change;
        double iteration_time;
        bool  converged;
        clock_t t1_inner, t2_inner;

        if (verbose) {
            std::cout << std::setw(7) << "Iter No" << "    "
                      << std::setw(8) << "Time(s)" << "    "
                      << std::setw(16) << "Change" << std::endl
                      << "-------    --------    ----------------\n";
        }

        t1 = clock();
        do {

            t1_inner = clock();
            sae.update(model);
            t2_inner = clock();
            converged = sae.converged(change);

            if (verbose) {
                std::cout << std::setw(7) << num_iterations + 1 << "    "
                          << std::setw(8) << std::setprecision(3)
                          << ((float)t2_inner - (float)t1_inner)/CLOCKS_PER_SEC
                          << "    " << std::setw(16) << std::setprecision(1)
                          << std::fixed << change << std::endl;
            }
        } while (++num_iterations < max_iterations && !converged);
        t2 = clock();

        diff = ((float)t2 - (float)t1) / CLOCKS_PER_SEC;
        std::cout << "Time for " << num_iterations << " iterations = " << diff
                  << "s\n";
        std::cout << "Time per iteration " << std::setprecision(2) << diff/num_iterations << "s\n";
    }

    std::cout << "Saving results to " << output_filename << std::endl;
    sae.saveStackSums(output_filename);
    sae.updateNoApplyTL(model);
    sae.saveStackSums(output_filename_counts);
    std::cout << "Done.\n";

    return 0;
}

void print_help()
{
    std::string title = "EMASE-Zero Version " VERSION " Help";

    std::cout << std::endl << std::endl
              << title << std::endl
              << std::string(title.length(), '-') << std::endl << std::endl
              << "USAGE: emase-zero [options] <alignment_incidence_file>\n\n"
              << "INPUT: Binary Alignment Incidence file prepared with alntools\n"
              << "       (https://churchill-lab.github.io/alntools/)\n\n"
              << "OPTIONS\n"
              << "  --model (-m) <int>:\n"
              << "      Specify normalization model (can be 1-4, default=1)\n\n"
              << "  --bin (-b):\n"
              << "      Binary input mode, with this option emase2 will expect a binary\n"
              << "      file exported from Kallisto as input\n\n"
              << "  --output (-o) <filename>:\n"
              << "      Specify filename for output file (default is input_basename.stacksum.tsv)\n\n"
              << "  --transcript-lengths (-l) <filename>:\n"
              << "      Filename for transcript length file. Format of each line is\n"
              << "     \"TranscriptName_HaplotypeName\\tlength\"\n\n"
              << "  --gene-mappings (-g) <filename>:\n"
              << "      Filename containing transcript to gene mapping. Tab delimited file where\n"
              << "      the first field is the gene ID, all other fields are the transcript IDs\n"
              << "      that belong to that gene\n\n"
              << "  --max-iterations (-i) <int>:\n"
              << "      Specify the maximum number of EM iterations. (Default 999)\n\n"
              << "  --tolerance (-t) <double>:\n"
              << "      Specify the convergence threshold. emase will terminate when the\n"
              << "      sum of the aboslute value of differences in the stack sum from one\n"
              << "      iteration to the next is lower than this value multiplied by\n"
              << "      1,000,000 if performing length adjustment and multipled by #reads if\n"
              << "      not (Default = 0.0001)\n\n"
              << "  --verbose (-v):\n"
              <<"       Run in verbose mode\n\n"
              << "  --version:\n"
              << "      Print the version and exit\n\n"
              << "  --help (-h):\n"
              << "      Print this message and exit\n\n";
}
<commit_msg>[kbchoi] Fixing segfault<commit_after>/*
 * Copyright (c) 2015 The Jackson Laboratory
 *
 * This software was developed by Gary Churchill's Lab at The Jackson
 * Laboratory (see http://research.jax.org/faculty/churchill).
 *
 * This is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This software 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 General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this software. If not, see <http://www.gnu.org/licenses/>.
 */

//
//  emase.cpp
//
//
//  Created by Glen Beane on 8/20/14.
//
//

#include <iostream>
#include <iomanip>
#include <cmath>
#include <getopt.h>

#include "alignment_incidence_matrix.h"
#include "sample_allelic_expression.h"
#include "python_interface.h"
#include "kallisto_import.h"

#define VERSION "0.2.4"

void print_help();

int main(int argc, char **argv)
{
    AlignmentIncidenceMatrix *aim;

    bool binary_input = false;
    int verbose = 0;

    int num_iterations;
    int max_iterations = 999;

    SampleAllelicExpression::model model = SampleAllelicExpression::MODEL_1;
    int m;

    clock_t t1, t2;
    float diff;

    std::string input_filename;
    std::string output_filename;
    std::string output_filename_counts;
    std::string transcript_length_file;
    std::string extension = ".pcl.bz2";
    std::string gene_file;

    //getopt related variables
    int c;
    int option_index = 0;
    bool bad_args = false;

    double tolerance = 0.0001;

    static struct option long_options[] = {
        {"help", no_argument, 0, 'h'},
        {"model", required_argument, 0, 'm'},
        {"outbase", required_argument, 0, 'o'},
        {"transcript-lengths", required_argument, 0, 'l'},
        {"max-iterations", required_argument, 0, 'i'},
        {"gene-mappings", required_argument, 0, 'g'},
        {"verbose", no_argument, &verbose, 1},
        {"version", no_argument, 0, 'V'},
        {"tolerance", required_argument, 0, 't'},
        {0, 0, 0, 0}
    };

    while ((c = getopt_long(argc, argv, "hm:o:l:i:bg:vVc:t:", long_options,
                            &option_index)) != -1) {
        switch (c) {
            case 'h':
                print_help();
                return 0;

            case 'm':
                m = std::stoi(optarg);
                if (m < 1 || m > 4) {
                    std::cerr << "[ERROR] Invalid model number specified. Valid options: 1, 2, 3, 4\n";
                }

                switch (m) {
                    case 1:
                        model = SampleAllelicExpression::MODEL_1;
                        break;
                    case 2:
                        model = SampleAllelicExpression::MODEL_2;
                        break;
                    case 3:
                        model = SampleAllelicExpression::MODEL_3;
                        break;
                    case 4:
                        model = SampleAllelicExpression::MODEL_4;
                        break;
                }

                break;

            case 'o':
                output_filename = std::string(optarg);
                output_filename_counts = std::string(optarg);
                output_filename.append(".tpm");
                output_filename_counts.append(".counts");
                break;

            case 'l':
                transcript_length_file = std::string(optarg);
                break;

            case 'i':
                max_iterations = std::stoi(optarg);
                break;

            case 'g':
                gene_file = std::string(optarg);
                break;

            case 'v':
                verbose = 1;
                break;

            case 'V':
               std::cout << VERSION << std::endl;
               return 0;

            case 't':
                tolerance = std::stod(optarg);
                break;

            case '?':
                bad_args = true;

        }
    }

    if (bad_args) {
        print_help();
        return 1;
    }

    if (argc - optind == 1) {
        input_filename = argv[optind];
    } else {
        std::cerr << "\n[ERROR] Missing required argument (input file name)\n";
        print_help();
        return 1;
    }

    if (output_filename.empty()) {
        //use default, based on the input file name but placed in current working directdory
        output_filename = "emase-zero.quantified.tpm";
        output_filename_counts = "emase-zero.quantified.counts";
    }

    std::cout << "\nemase-zero Version " << VERSION << std::endl <<std::endl;

    std::cout << "Alignment File: " << input_filename << std::endl;

    if (!gene_file.empty()) {
        std::cout << "Grouping File: " << gene_file << std::endl;
    }
    else {
        std::cout << "Grouping File: None\n";
    }

    if (!transcript_length_file.empty()) {
        std::cout << "Transcript Length File: " << transcript_length_file
                  << std::endl;
    }
    else {
        std::cout << "Transcript Length File: None\n";
    }

    std::cout << "EM Model: " << model << std::endl
              << "Output File (TPM): " << output_filename << std::endl
              << "Output File (Counts): " << output_filename_counts << std::endl
              << "----------------------------------------------------\n\n\n";

    std::cout << "Loading " << input_filename << "..." << std::endl;
    aim = loadFromBin(input_filename);

    if (!aim) {
        std::cerr << "Error loading binary input file\n";
        return 1;
    }

    std::vector<std::string> hap_names = aim->get_haplotype_names();

    if (verbose) {
        std::cout << "File had the following haplotype names:\n";
        for (auto it = hap_names.begin(); it != hap_names.end(); ++it) {
            std::cout << *it << "\t";
        }
        std::cout << std::endl;

        if (aim->has_equivalence_classes()) {
            std::cout << aim->num_alignment_classes()
                      << " alignment classes loaded ("
                      << aim->total_reads() << " total reads)\n";
        }
        else {
            std::cout << aim->total_reads() << " reads loaded\n";
        }
        std::cout << aim->num_transcripts() << " transcripts\n";

        std::cout << std::endl;
    }

    if (!transcript_length_file.empty()) {
        if (verbose) {
            std::cout << "Loading Transcript Length File "
                      << transcript_length_file << std::endl;
        }
        aim->loadTranscriptLengths(transcript_length_file);
    }

    if (!gene_file.empty()) {
        if (verbose) {
            std::cout << "Loading Gene Mapping File " << gene_file << std::endl;
        }
        aim->loadGeneMappings(gene_file);
    }

    if (model != SampleAllelicExpression::MODEL_4 && !aim->has_gene_mappings()) {
        if (!binary_input) {
            std::cerr << "[ERROR] File does not contain transcript to gene mapping information.  Only normalization Model 4 can be used.\n";
        }
        else {
            std::cerr << "[ERROR] Only model 4 is possible without gene to transcript information (--gene-mappings, -g). You specified another model.\n";
        }
        return 1;
    }

    t1 = clock();
    SampleAllelicExpression sae(aim, tolerance);
    t2 = clock();

    if (verbose) {
        diff = ((float)t2-(float)t1)/CLOCKS_PER_SEC;
        std::cout << "Time for initializing stack sum = " << diff << "s"
                  << std::endl;
    }

    if (max_iterations > 0) {
        num_iterations = 0;
        std::cout << "Beginning EM Iterations..." << std::endl;
        double change;
        double iteration_time;
        bool  converged;
        clock_t t1_inner, t2_inner;

        if (verbose) {
            std::cout << std::setw(7) << "Iter No" << "    "
                      << std::setw(8) << "Time(s)" << "    "
                      << std::setw(16) << "Change" << std::endl
                      << "-------    --------    ----------------\n";
        }

        t1 = clock();
        do {

            t1_inner = clock();
            sae.update(model);
            t2_inner = clock();
            converged = sae.converged(change);

            if (verbose) {
                std::cout << std::setw(7) << num_iterations + 1 << "    "
                          << std::setw(8) << std::setprecision(3)
                          << ((float)t2_inner - (float)t1_inner)/CLOCKS_PER_SEC
                          << "    " << std::setw(16) << std::setprecision(1)
                          << std::fixed << change << std::endl;
            }
        } while (++num_iterations < max_iterations && !converged);
        t2 = clock();

        diff = ((float)t2 - (float)t1) / CLOCKS_PER_SEC;
        std::cout << "Time for " << num_iterations << " iterations = " << diff
                  << "s\n";
        std::cout << "Time per iteration " << std::setprecision(2) << diff/num_iterations << "s\n";
    }

    std::cout << "Saving results to " << output_filename << std::endl;
    sae.saveStackSums(output_filename);
    sae.updateNoApplyTL(model);
    sae.saveStackSums(output_filename_counts);
    std::cout << "Done.\n";

    return 0;
}

void print_help()
{
    std::string title = "EMASE-Zero Version " VERSION " Help";

    std::cout << std::endl << std::endl
              << title << std::endl
              << std::string(title.length(), '-') << std::endl << std::endl
              << "USAGE: emase-zero [options] <alignment_incidence_file>\n\n"
              << "INPUT: Binary Alignment Incidence file prepared with alntools\n"
              << "       (https://churchill-lab.github.io/alntools/)\n\n"
              << "OPTIONS\n"
              << "  --model (-m) <int>:\n"
              << "      Specify normalization model (can be 1-4, default=1)\n\n"
              << "  --bin (-b):\n"
              << "      Binary input mode, with this option emase2 will expect a binary\n"
              << "      file exported from Kallisto as input\n\n"
              << "  --output (-o) <filename>:\n"
              << "      Specify filename for output file (default is input_basename.stacksum.tsv)\n\n"
              << "  --transcript-lengths (-l) <filename>:\n"
              << "      Filename for transcript length file. Format of each line is\n"
              << "     \"TranscriptName_HaplotypeName\\tlength\"\n\n"
              << "  --gene-mappings (-g) <filename>:\n"
              << "      Filename containing transcript to gene mapping. Tab delimited file where\n"
              << "      the first field is the gene ID, all other fields are the transcript IDs\n"
              << "      that belong to that gene\n\n"
              << "  --max-iterations (-i) <int>:\n"
              << "      Specify the maximum number of EM iterations. (Default 999)\n\n"
              << "  --tolerance (-t) <double>:\n"
              << "      Specify the convergence threshold. emase will terminate when the\n"
              << "      sum of the aboslute value of differences in the stack sum from one\n"
              << "      iteration to the next is lower than this value multiplied by\n"
              << "      1,000,000 if performing length adjustment and multipled by #reads if\n"
              << "      not (Default = 0.0001)\n\n"
              << "  --verbose (-v):\n"
              <<"       Run in verbose mode\n\n"
              << "  --version:\n"
              << "      Print the version and exit\n\n"
              << "  --help (-h):\n"
              << "      Print this message and exit\n\n";
}
<|endoftext|>


<commit_before>/***************************************************************************
 *  @file       RedBlackTree.hpp
 *  @author     Yue Wang
 *  @date       02  July 2014
 *  @remark     Algorithims implementation for CLRS, using C++ templates.
 ***************************************************************************/

#ifndef REDBLACKTREE_H
#define REDBLACKTREE_H

#include "Node.hpp"

namespace ch13{

template<typename K, typename D>
class RedBlackTree
{
public:
    using KeyType   =   K;
    using DataType  =   D;
    using sPointer  =   std::shared_ptr<Node<K,D> >;
    using wPointer  =   std::weak_ptr<Node<K,D> >;

    /**
     * @brief default ctor
     */
    RedBlackTree():
        root(),nil(std::make_shared<Node<K,D>>())
    {}

private:
    sPointer root;
    sPointer nil;

    /**
     * @brief left_rotate
     *
     * @complexity  O(1)
     *
     * based on LEFT-ROTATE     Page 313
     *      x   :   target
     *      y   :   new_parent
     */
    void left_rotate(sPointer target)
    {
        sPointer new_parent = target->right;

        //! turn new_parent's left subtree into target's right subtree
        target->right   =   new_parent->left;
        if(target->left != this->nil)
            new_parent->left->parent.lock()  =  target;

        //! link target's parent to new_parent
        new_parent->parent  =   target->parent;
        if(target->parent   ==  this->nil)
            this->root  =   new_parent;
        else
        {
            sPointer pnt = target->parent.lock();
            (target == pnt->left?   pnt->left   :   pnt->right)    =   new_parent;
        }

        //! put target on new_parent’s left
        new_parent->left    =   target;
        target->parent      =   new_parent;
    }

    /**
     * @brief right_rotate
     *
     * @complexity  O(1)
     *
     * as required for ex13.2-1
     */
    void righ_rotate(sPointer target)
    {
        sPointer new_parent = target->left;

        //! turn new_parent's right subtree into target's left subtree
        target->left    =   new_parent->right;
        if(new_parent->right    !=  this->nil)
            new_parent->right->parent.lock()    =   target;

        //! link target's parent to new_parent
        if(target->parent.lock()    ==  this->nil)
            this->root  =   new_parent;
        else
        {
            sPointer pnt = target->parent.lock();
            (target == pnt->left?   pnt->left   :   pnt->right)     =   new_parent;
        }

        //! put target on new_parent’s right
        new_parent->right   =   target;
        target->parent      =   new_parent;
    }

    /**
     * @brief insert_fixup
     *
     * @complexity  O(1)
     *
     * based on that on Page 316
     */
    void insert_fixup(sPointer target)
    {
        while(target->parent.lock()->color ==  Color::Red)
        {
            if(target->parent.lock()    ==
               target->parent.lock()->parent.lock()->left)
            {
                sPointer uncle  =   target->parent.lock()->parent.lock()->right;
                if(uncle->color ==  Color::Red)
                {
                    uncle->color = Color::Black;
                    target->parent.lock()->color = Color::Black;
                    target->parent.lock()->parent.lock()->color = Color::Red;
                }
                else
                {
                    if(target   ==  target->parent.lock()->right)
                    {
                        target  =   target->parent.lock();
                        left_rotate(target);
                    }
                    target->parent.lock()->color                =   Color::Black;
                    target->parent.lock()->parent.lock()->color =   Color::Red;
                    righ_rotate(target->parent.lock()->parent.lock());
                }
            }
            else
            {
                sPointer uncle  =   target->parent.lock()->parent.lock()->left;
                if(uncle->color ==  Color::Red)
                {
                    uncle->color = Color::Black;
                    target->parent.lock()->color = Color::Black;
                    target->parent.lock()->parent.lock()->color = Color::Red;
                }
                else
                {
                    if(target   ==  target->parent.lock()->left)
                    {
                        target  =   target->parent.lock();
                        right_rotate(target);
                    }
                    target->parent.lock()->color                =   Color::Black;
                    target->parent.lock()->parent.lock()->color =   Color::Red;
                    left_rotate(target->parent.lock()->parent.lock());
                }
            }
        }
    }
};

}//namespace ch13
#endif // REDBLACKTREE_H
<commit_msg>added insert() not tested yet  On branch ch13 	modified:   ch13/RedBlackTree.hpp<commit_after>/***************************************************************************
 *  @file       RedBlackTree.hpp
 *  @author     Yue Wang
 *  @date       02  July 2014
 *  @remark     Algorithims implementation for CLRS, using C++ templates.
 ***************************************************************************/

#ifndef REDBLACKTREE_H
#define REDBLACKTREE_H

#include "Node.hpp"

namespace ch13{

/**
 * @brief Red Black Tree
 */
template<typename K, typename D>
class RedBlackTree
{
public:
    using KeyType   =   K;
    using DataType  =   D;
    using sPointer  =   std::shared_ptr<Node<K,D> >;
    using wPointer  =   std::weak_ptr<Node<K,D> >;

    /**
     * @brief default ctor
     */
    RedBlackTree():
        root(),nil(std::make_shared<Node<K,D>>())
    {}

    /**
     * @brief insert
     *
     * @complexity  O(h)
     *
     * based on that of Page 316
     */
    void insert(sPointer added)
    {
        sPointer tracker    =   this->nil;
        sPointer curr       =   this->root;
        while(curr != this->nil)
        {
            tracker =   curr;
            curr    =   (curr->key > added->key?    curr->left  :   curr->right);
        }

        added->parent   =   tracker;
        if(tracker  ==  this->nil)
            this->root  =   added;
        else
            (added->key < tracker->key?     tracker->left   :   tracker->right)
                        =   added;

        added->left     =   added->right    =   this->nil;
        added->color    =   Color::Red;
        insert_fixup(added);
    }

private:
    sPointer root;
    sPointer nil;

    /**
     * @brief left_rotate
     *
     * @complexity  O(1)
     *
     * based on LEFT-ROTATE     Page 313
     *      x   :   target
     *      y   :   new_parent
     */
    void left_rotate(sPointer target)
    {
        sPointer new_parent = target->right;

        //! turn new_parent's left subtree into target's right subtree
        target->right   =   new_parent->left;
        if(target->left != this->nil)
            new_parent->left->parent.lock()  =  target;

        //! link target's parent to new_parent
        new_parent->parent  =   target->parent;
        if(target->parent   ==  this->nil)
            this->root  =   new_parent;
        else
        {
            sPointer pnt = target->parent.lock();
            (target == pnt->left?   pnt->left   :   pnt->right)    =   new_parent;
        }

        //! put target on new_parent’s left
        new_parent->left    =   target;
        target->parent      =   new_parent;
    }

    /**
     * @brief right_rotate
     *
     * @complexity  O(1)
     *
     * as required for ex13.2-1
     */
    void righ_rotate(sPointer target)
    {
        sPointer new_parent = target->left;

        //! turn new_parent's right subtree into target's left subtree
        target->left    =   new_parent->right;
        if(new_parent->right    !=  this->nil)
            new_parent->right->parent.lock()    =   target;

        //! link target's parent to new_parent
        if(target->parent.lock()    ==  this->nil)
            this->root  =   new_parent;
        else
        {
            sPointer pnt = target->parent.lock();
            (target == pnt->left?   pnt->left   :   pnt->right)     =   new_parent;
        }

        //! put target on new_parent’s right
        new_parent->right   =   target;
        target->parent      =   new_parent;
    }

    /**
     * @brief insert_fixup
     *
     * @complexity  O(1)
     *
     * based on that on Page 316
     */
    void insert_fixup(sPointer target)
    {
        while(target->parent.lock()->color ==  Color::Red)
        {
            if(target->parent.lock()    ==
               target->parent.lock()->parent.lock()->left)
            {
                sPointer uncle  =   target->parent.lock()->parent.lock()->right;
                if(uncle->color ==  Color::Red)
                {
                    uncle->color = Color::Black;
                    target->parent.lock()->color = Color::Black;
                    target->parent.lock()->parent.lock()->color = Color::Red;
                }
                else
                {
                    if(target   ==  target->parent.lock()->right)
                    {
                        target  =   target->parent.lock();
                        left_rotate(target);
                    }
                    target->parent.lock()->color                =   Color::Black;
                    target->parent.lock()->parent.lock()->color =   Color::Red;
                    righ_rotate(target->parent.lock()->parent.lock());
                }
            }
            else
            {
                sPointer uncle  =   target->parent.lock()->parent.lock()->left;
                if(uncle->color ==  Color::Red)
                {
                    uncle->color = Color::Black;
                    target->parent.lock()->color = Color::Black;
                    target->parent.lock()->parent.lock()->color = Color::Red;
                }
                else
                {
                    if(target   ==  target->parent.lock()->left)
                    {
                        target  =   target->parent.lock();
                        right_rotate(target);
                    }
                    target->parent.lock()->color                =   Color::Black;
                    target->parent.lock()->parent.lock()->color =   Color::Red;
                    left_rotate(target->parent.lock()->parent.lock());
                }
            }
        }
    }
};

}//namespace ch13
#endif // REDBLACKTREE_H
<|endoftext|>


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

  Program:   Insight Segmentation & Registration Toolkit (ITK)
  Module:    itkPointSetToImageTranslationNormalizedCorrelationRegularStepGradientDescentRegistrationTest.cxx
  Language:  C++
  Date:      $Date$
  Version:   $Revision$
s

Copyright (c) 2000 National Library of Medicine
All rights reserved.

See COPYRIGHT.txt for copyright details.

=========================================================================*/
#include "itkPhysicalImage.h"
#include "itkSimpleImageRegionIterator.h"
#include "itkPointSetToImageTranslationNormalizedCorrelationRegularStepGradientDescentRegistration.h"

/**
 *  This test uses two 2D-Gaussians (standard deviation RegionSize/2)
 *
 *  One of the images is subsampled in order to obtain a PoinSet.
 *  The PointSet is considered the Target and the other image is
 *  considered the Reference.
 */

int main()
{

  // Number of pixels to take from the target image
  unsigned long  numberOfSamples = 10000L;

  /*Allocate Images*/
  const unsigned int ImageDimension  = 2;

  typedef itk::PhysicalImage<unsigned char,ImageDimension>  ReferenceType;

  typedef itk::DefaultStaticMeshTraits<unsigned char,   // pixel type
                                       ImageDimension,  // point dimension
                                       1,               // max topological dim
                                       double,          // coordinates type
                                       double           // interpolation weight
                                         >  DefaultPointSetTraits;

  typedef itk::PointSet<unsigned char, ImageDimension, DefaultPointSetTraits>    PointSetType;

  typedef PointSetType                                          TargetType;

  typedef itk::PointSetToImageTranslationNormalizedCorrelationRegularStepGradientDescentRegistration<
                                                                ReferenceType,
                                                                TargetType> RegistrationType;

  ReferenceType::SizeType size = {{100,100}};
  ReferenceType::IndexType index = {{0,0}};
  ReferenceType::RegionType region;
  region.SetSize( size );
  region.SetIndex( index );

  ReferenceType::Pointer imgReference = ReferenceType::New();
  imgReference->SetLargestPossibleRegion( region );
  imgReference->SetBufferedRegion( region );
  imgReference->SetRequestedRegion( region );
  imgReference->Allocate();

  ReferenceType::Pointer imgTarget = ReferenceType::New(); // type=Reference but this
                                                           // is the target image to subsample
  imgTarget->SetLargestPossibleRegion( region );
  imgTarget->SetBufferedRegion( region );
  imgTarget->SetRequestedRegion( region );
  imgTarget->Allocate();

  /* Fill images with a 2D gaussian*/
  typedef  itk::SimpleImageRegionIterator<ReferenceType> ReferenceIteratorType;

  itk::Point<double,2> center;
  center[0] = (double)region.GetSize()[0]/2.0;
  center[1] = (double)region.GetSize()[1]/2.0;

  const double s = (double)region.GetSize()[0]/2.0;

  itk::Point<double,2>  p;
  itk::Vector<double,2> d;

  /* Set the displacement */
  itk::Vector<double,2> displacement;
  displacement[0] = 7;
  displacement[1] =  3;

  ReferenceIteratorType ri(imgReference,region);
  ReferenceIteratorType ti(imgTarget,region);
  ri.Begin();
  while(!ri.IsAtEnd())
  {
    p[0] = ri.GetIndex()[0];
    p[1] = ri.GetIndex()[1];
    d = p-center;
    d += displacement;
    const double x = d[0];
    const double y = d[1];
    ri.Set( 200.0 * exp( - ( x*x + y*y )/(s*s) ) );
    ++ri;
  }

  ti.Begin();
  while(!ti.IsAtEnd())
  {
    p[0] = ti.GetIndex()[0];
    p[1] = ti.GetIndex()[1];
    d = p-center;
    const double x = d[0];
    const double y = d[1];
    ti.Set( 200.0 * exp( - ( x*x + y*y )/(s*s) ) );
    ++ti;
  }

  // Subsample the target image to produce a point set
  PointSetType::Pointer pointSetTarget = PointSetType::New();

  PointSetType::PointsContainer::Pointer    points = pointSetTarget->GetPoints();
  PointSetType::PointDataContainer::Pointer data   = PointSetType::PointDataContainer::New();
  pointSetTarget->SetPointData( data );

  points->Reserve( numberOfSamples );
  data->Reserve(   numberOfSamples );

  const unsigned int numPixelsToSkip =
            imgTarget->GetOffsetTable()[ImageDimension] / numberOfSamples;

  unsigned int counter   = 0;
  unsigned int numPoints = 0;

  ti.Begin();
  while(!ti.IsAtEnd() && numPoints < numberOfSamples )
  {
    if( counter < numPixelsToSkip )
    {
      PointSetType::PointType point;
      point[0] = ti.GetIndex()[0];
      point[1] = ti.GetIndex()[1];
      points->SetElement( counter, point    );
      data->SetElement(   counter, ti.Get() );
      numPoints++;
    }
    ++ti;
    ++counter;
  }

  RegistrationType::Pointer registrationMethod = RegistrationType::New();

  registrationMethod->SetReference(imgReference);
  registrationMethod->SetTarget(pointSetTarget);

  registrationMethod->StartRegistration();

  std::cout << "The correct answer should be : " << std::endl;
  std::cout << -displacement << std::endl;

  return EXIT_SUCCESS;

}

<commit_msg>BUG: counter and numPoints were mixed up during point list initialization<commit_after>/*=========================================================================

  Program:   Insight Segmentation & Registration Toolkit (ITK)
  Module:    itkPointSetToImageTranslationNormalizedCorrelationRegularStepGradientDescentRegistrationTest.cxx
  Language:  C++
  Date:      $Date$
  Version:   $Revision$
s

Copyright (c) 2000 National Library of Medicine
All rights reserved.

See COPYRIGHT.txt for copyright details.

=========================================================================*/
#include "itkPhysicalImage.h"
#include "itkSimpleImageRegionIterator.h"
#include "itkPointSetToImageTranslationNormalizedCorrelationRegularStepGradientDescentRegistration.h"

/**
 *  This test uses two 2D-Gaussians (standard deviation RegionSize/2)
 *
 *  One of the images is subsampled in order to obtain a PoinSet.
 *  The PointSet is considered the Target and the other image is
 *  considered the Reference.
 */

int main()
{

  // Number of pixels to take from the target image
  unsigned long  numberOfSamples = 10000L;

  /*Allocate Images*/
  const unsigned int ImageDimension  = 2;

  typedef itk::PhysicalImage<unsigned char,ImageDimension>  ReferenceType;

  typedef itk::DefaultStaticMeshTraits<unsigned char,   // pixel type
                                       ImageDimension,  // point dimension
                                       1,               // max topological dim
                                       double,          // coordinates type
                                       double           // interpolation weight
                                         >  DefaultPointSetTraits;

  typedef itk::PointSet<unsigned char, ImageDimension, DefaultPointSetTraits>    PointSetType;

  typedef PointSetType                                          TargetType;

  typedef itk::PointSetToImageTranslationNormalizedCorrelationRegularStepGradientDescentRegistration<
                                                                ReferenceType,
                                                                TargetType> RegistrationType;

  ReferenceType::SizeType size = {{100,100}};
  ReferenceType::IndexType index = {{0,0}};
  ReferenceType::RegionType region;
  region.SetSize( size );
  region.SetIndex( index );

  ReferenceType::Pointer imgReference = ReferenceType::New();
  imgReference->SetLargestPossibleRegion( region );
  imgReference->SetBufferedRegion( region );
  imgReference->SetRequestedRegion( region );
  imgReference->Allocate();

  ReferenceType::Pointer imgTarget = ReferenceType::New(); // type=Reference but this
                                                           // is the target image to subsample
  imgTarget->SetLargestPossibleRegion( region );
  imgTarget->SetBufferedRegion( region );
  imgTarget->SetRequestedRegion( region );
  imgTarget->Allocate();

  /* Fill images with a 2D gaussian*/
  typedef  itk::SimpleImageRegionIterator<ReferenceType> ReferenceIteratorType;

  itk::Point<double,2> center;
  center[0] = (double)region.GetSize()[0]/2.0;
  center[1] = (double)region.GetSize()[1]/2.0;

  const double s = (double)region.GetSize()[0]/2.0;

  itk::Point<double,2>  p;
  itk::Vector<double,2> d;

  /* Set the displacement */
  itk::Vector<double,2> displacement;
  displacement[0] = 7;
  displacement[1] =  3;

  ReferenceIteratorType ri(imgReference,region);
  ReferenceIteratorType ti(imgTarget,region);
  ri.Begin();
  while(!ri.IsAtEnd())
  {
    p[0] = ri.GetIndex()[0];
    p[1] = ri.GetIndex()[1];
    d = p-center;
    d += displacement;
    const double x = d[0];
    const double y = d[1];
    ri.Set( 200.0 * exp( - ( x*x + y*y )/(s*s) ) );
    ++ri;
  }

  ti.Begin();
  while(!ti.IsAtEnd())
  {
    p[0] = ti.GetIndex()[0];
    p[1] = ti.GetIndex()[1];
    d = p-center;
    const double x = d[0];
    const double y = d[1];
    ti.Set( 200.0 * exp( - ( x*x + y*y )/(s*s) ) );
    ++ti;
  }

  // Subsample the target image to produce a point set
  PointSetType::Pointer pointSetTarget = PointSetType::New();

  PointSetType::PointsContainer::Pointer    points = pointSetTarget->GetPoints();
  PointSetType::PointDataContainer::Pointer data   = PointSetType::PointDataContainer::New();
  pointSetTarget->SetPointData( data );

  points->Reserve( numberOfSamples );
  data->Reserve(   numberOfSamples );

  const unsigned int numPixelsToSkip =
            imgTarget->GetOffsetTable()[ImageDimension] / numberOfSamples;

  unsigned int counter   = 0;
  unsigned int numPoints = 0;

  ti.Begin();
  while(!ti.IsAtEnd() && numPoints < numberOfSamples )
  {
    if( counter >= numPixelsToSkip )
    {
      PointSetType::PointType point;
      point[0] = ti.GetIndex()[0];
      point[1] = ti.GetIndex()[1];
      points->SetElement( numPoints, point    );
      data->SetElement(   numPoints, ti.Get() );
      numPoints++;
      counter = 0;
    }
    ++ti;
    ++counter;
  }

  RegistrationType::Pointer registrationMethod = RegistrationType::New();

  registrationMethod->SetReference(imgReference);
  registrationMethod->SetTarget(pointSetTarget);

  registrationMethod->StartRegistration();

  std::cout << std::endl << "After  " <<
    registrationMethod->GetOptimizer()->GetCurrentNumberOfIterations()
    << "  Iterations " << std::endl;

  // get the results
  RegistrationType::ParametersType solution =
    registrationMethod->GetOptimizer()->GetCurrentPosition();

  std::cout << "Solution is: " << solution << std::endl;

  //
  // check results to see if it is within range
  //
  bool pass = true;
  double trueParameters[2] = { -7, -3 };
  for( unsigned int j = 0; j < 2; j++ )
    {
    if( vnl_math_abs( solution[j] - trueParameters[j] ) > 0.02 )
      pass = false;
    }

  if( !pass )
    {
    std::cout << "Test failed." << std::endl;
    return EXIT_FAILURE;
    }

  std::cout << "Test passed." << std::endl;
  return EXIT_SUCCESS;

}

<|endoftext|>