dsslib/sort/sort.c

/***********************************************************************
*                                                                      *
*               This software is part of the ast package               *
*          Copyright (c) 2011-2012 AT&T Intellectual Property          *
*                      and is licensed under the                       *
*                 Eclipse Public License, Version 1.0                  *
*                    by AT&T Intellectual Property                     *
*                                                                      *
*                A copy of the License is available at                 *
*          http://www.eclipse.org/org/documents/epl-v10.html           *
*         (with md5 checksum b35adb5213ca9657e911e9befb180842)         *
*                                                                      *
*              Information and Software Systems Research               *
*                            AT&T Research                             *
*                           Florham Park NJ                            *
*                                                                      *
*                 Glenn Fowler <gsf@research.att.com>                  *
*                                                                      *
***********************************************************************/
#pragma prototyped

static const char sort_usage[] =
"[+PLUGIN?\findex\f]"
"[+DESCRIPTION?The sort query writes the sorted input records to the"
"   standard output. The unsorted record stream is passed to the"
"   next query, if any. The sort keys are the \afield\a operands,"
"   or the raw record data if there are no operands.]"
"[c:count?Prepend an integer count field of the number of records that "
    "compare equal. Records with count less that \athreshhold\a are "
    "omitted.]#?[threshold:=1]"
"[r:reverse|invert?Reverse the sense of comparisons.]"
"[u:unique?Keep only the first of multiple records that compare equal on "
    "all keys.]"
"\n"
"\n[ field ... ]\n"
"\n"
"[+CAVEATS?Currently all data is sorted in memory -- spillover to "
    "temporary files not implemented yet.]"
;

#include <dsslib.h>
#include <recsort.h>
#include <stk.h>

struct State_s; typedef struct State_s State_t;

struct State_s
{
    Rsdisc_t    sortdisc;
    Rsdisc_t    uniqdisc;
    Rskeydisc_t keydisc;
    Rskey_t*    sortkey;
    Rskey_t*    uniqkey;
    Rs_t*       sort;
    Rs_t*       uniq;
    Cx_t*       cx;
    Sfio_t*     op;
    Sfio_t*     sortstack;
    Sfio_t*     uniqstack;
    Sfio_t*     tmp;
    char*       sortbase;
    char*       uniqbase;
    void*       data;
    Dssfile_t*  file;
    Vmalloc_t*  rm;
    Vmalloc_t*  vm;
    size_t      count;
};

extern Dsslib_t dss_lib_sort;

static ssize_t
key(Rs_t* rs, unsigned char* data, size_t datasize, unsigned char* key, size_t keysize, Rsdisc_t* disc)
{
    State_t*    state = (State_t*)disc;
    Rskeyfield_t*   field;
    Cxoperand_t r;
    unsigned char*  k;
    unsigned char*  e;

    k = key;
    e = k + keysize;
    for (field = state->sortkey->head; field; field = field->next)
    {
        if (cxcast(state->cx, &r, (Cxvariable_t*)field->user, state->cx->state->type_string, state->data, NiL))
            return -1;
        k += field->coder(state->sortkey, field, (unsigned char*)r.value.string.data, r.value.string.size, k, e);
    }
    return k - key;
}

static ssize_t
rev(Rs_t* rs, unsigned char* data, size_t datasize, unsigned char* key, size_t keysize, Rsdisc_t* disc)
{
    State_t*    state = (State_t*)disc;

    return state->sortkey->head->coder(state->sortkey, state->sortkey->head, data, datasize, key, key + keysize);
}

static int
count(Rs_t* rs, int op, void* data, void* arg, Rsdisc_t* disc)
{
    State_t*    state = (State_t*)disc;
    Rsobj_t*    r;
    Rsobj_t*    q;
    char*       s;
    ssize_t     n;

    switch (op)
    {
    case RS_POP:
        break;
    case RS_WRITE:
        r = (Rsobj_t*)data;
        n = 1;
        for (q = r->equal; q; q = q->right)
            n++;
        if (n >= state->count)
        {
            n = sfprintf(state->uniqstack, "%I*u %-.*s", sizeof(n), n, r->datalen, r->data);
            s = stkfreeze(state->uniqstack, 0);
            if (rsprocess(state->uniq, s, -n) <= 0)
            {
                if (state->cx->disc->errorf)
                    (*state->cx->disc->errorf)(state->cx, disc, ERROR_SYSTEM|2, "uniq record process error");
                return -1;
            }
        }
        return RS_DELETE;
    default:
        return -1;
    }
    return 0;
}

static int
sort_beg(Cx_t* cx, Cxexpr_t* expr, void* data, Cxdisc_t* disc)
{
    char**      argv = (char**)data;
    int     errors = error_info.errors;
    int     n;
    int     uniq;
    char*       s;
    char*       t;
    char*       num;
    State_t*    state;
    Cxvariable_t*   variable;
    Vmalloc_t*  vm;
    char        opt[2];

    if (!(vm = vmopen(Vmdcheap, Vmlast, 0)) || !(state = vmnewof(vm, 0, State_t, 1, 0)))
    {
        if (vm)
            vmclose(vm);
        if (disc->errorf)
            (*disc->errorf)(cx, disc, ERROR_SYSTEM|2, "out of space");
        return -1;
    }
    state->vm = vm;
    state->keydisc.version = RSKEY_VERSION;
    state->keydisc.errorf = disc->errorf;
    if (!(state->sortkey = rskeyopen(&state->keydisc, &state->sortdisc)))
        goto bad;
    if (!(state->sort = rsnew(state->sortkey->disc)))
        goto bad;
    if (!(state->sortstack = stkopen(0)))
        goto bad;
    if (!(state->tmp = sfstropen()))
        goto bad;
    state->sortbase = stkptr(state->sortstack, 0);
    if (!(state->file = dssfopen(DSS(cx), "-", state->sortstack, DSS_FILE_WRITE, 0)))
        goto bad;
    sfprintf(cx->buf, "%s%s", strchr(dss_lib_sort.description, '['), sort_usage);
    s = sfstruse(cx->buf);
    uniq = 0;
    for (;;)
    {
        switch (optget(argv, s))
        {
        case 0:
            break;
        case 'c':
            state->count = (size_t)opt_info.number;
            continue;
        case 'u':
            uniq = 1;
            continue;
        case '?':
            if (disc->errorf)
            {
                (*disc->errorf)(cx, disc, ERROR_USAGE|4, "%s", opt_info.arg);
            }
            else
                goto bad;
            continue;
        case ':':
            if (disc->errorf)
                (*disc->errorf)(cx, disc, 2, "%s", opt_info.arg);
            else
                goto bad;
            continue;
        default:
            opt[0] = opt_info.option[1];
            opt[1] = 0;
            if (rskeyopt(state->sortkey, opt, 1))
                goto bad;
            continue;
        }
        break;
    }
    if (error_info.errors > errors)
        goto bad;
    argv += opt_info.index;
    n = 0;
    num = state->sortkey->head && state->sortkey->head->rflag ? "nr" : "n";
    while (s = *argv++)
    {
        if (t = strchr(s, '-'))
        {
            sfwrite(cx->buf, s, t - s);
            s = sfstruse(cx->buf);
            t++;
        }
        else
            t = 0;
        if (!(variable = cxvariable(cx, s, NiL, disc)))
            goto bad;
        if (rskey(state->sortkey, t ? t : cxisnumber(variable->type) ? num : "", 0))
            goto bad;
        state->sortkey->tail->user = variable;
        n = 1;
    }
    if (uniq)
    {
        state->sortkey->type &= ~RS_DATA;
        state->sortkey->type |= RS_UNIQ;
    }
    if (state->count)
    {
        state->sortdisc.events |= RS_WRITE;
        state->sortdisc.eventf = count;
        if (!(state->uniqstack = stkopen(0)))
            goto bad;
        state->uniqbase = stkptr(state->uniqstack, 0);
        if (!(state->uniqkey = rskeyopen(&state->keydisc, &state->uniqdisc)))
            goto bad;
        if (!(state->uniq = rsnew(state->uniqkey->disc)))
            goto bad;
        if (rskey(state->uniqkey, "1n", 0))
            goto bad;
        if (rskeyinit(state->uniqkey))
        {
            if (disc->errorf)
                (*disc->errorf)(cx, disc, ERROR_SYSTEM|2, "uniq key error");
            goto bad;
        }
        if (rsinit(state->uniq, state->uniqkey->meth, state->uniqkey->procsize, state->uniqkey->type, state->uniqkey))
        {
            if (disc->errorf)
                (*disc->errorf)(cx, disc, ERROR_SYSTEM|2, "uniq initialization error");
            goto bad;
        }
    }
    if (rskeyinit(state->sortkey))
    {
        if (disc->errorf)
            (*disc->errorf)(cx, disc, ERROR_SYSTEM|2, "sort key error");
        goto bad;
    }
    if (n)
    {
        state->sortdisc.defkeyf = key;
        state->sortdisc.key = 1;
    }
    else if (state->sortkey->head->rflag)
    {
        state->sortdisc.defkeyf = rev;
        state->sortdisc.key = 1;
    }
    if (rsinit(state->sort, state->sortkey->meth, state->sortkey->procsize, state->sortkey->type, state->sortkey))
    {
        if (disc->errorf)
            (*disc->errorf)(cx, disc, ERROR_SYSTEM|2, "sort initialization error");
        goto bad;
    }
    state->cx = cx;
    state->op = expr->op;
    expr->data = state;
    return 0;
 bad:
    if (state->sort)
        rsclose(state->sort);
    if (state->uniq)
        rsclose(state->uniq);
    if (state->sortkey)
        rskeyclose(state->sortkey);
    if (state->uniqkey)
        rskeyclose(state->uniqkey);
    if (state->file)
        dssfclose(state->file);
    else if (state->sortstack)
        stkclose(state->sortstack);
    if (state->uniqstack)
        stkclose(state->uniqstack);
    if (state->tmp)
        sfstrclose(state->tmp);
    vmclose(state->vm);
    return -1;
}

static int
sort_act(Cx_t* cx, Cxexpr_t* expr, void* data, Cxdisc_t* disc)
{
    State_t*    state = (State_t*)expr->data;
    Dssrecord_t*    record = (Dssrecord_t*)data;
    char*       s;
    ssize_t     n;

    if (dssfwrite(state->file, record))
        return -1;
    n = stktell(state->file->io);
    s = stkfreeze(state->file->io, 0);
    state->data = data;
    if (rsprocess(state->sort, s, -n) <= 0)
    {
        if (disc->errorf)
            (*disc->errorf)(cx, disc, ERROR_SYSTEM|2, "sort record process error");
        return -1;
    }
    return 0;
}

static int
sort_end(Cx_t* cx, Cxexpr_t* expr, void* data, Cxdisc_t* disc)
{
    State_t*    state = (State_t*)expr->data;
    int     r;

    r = 0;
    if (rswrite(state->sort, expr->op, RS_OTEXT))
    {
        if (disc->errorf)
            (*disc->errorf)(cx, disc, ERROR_SYSTEM|2, "sort output error");
        r = -1;
    }
    while (rsdisc(state->sort, NiL, RS_POP));
    if (rsclose(state->sort))
    {
        if (disc->errorf)
            (*disc->errorf)(cx, disc, ERROR_SYSTEM|2, "sort close error");
        r = -1;
    }
    if (rskeyclose(state->sortkey))
    {
        if (disc->errorf)
            (*disc->errorf)(cx, disc, ERROR_SYSTEM|2, "sort key error");
        r = -1;
    }
    if (state->uniq)
    {
        if (rswrite(state->uniq, expr->op, RS_OTEXT))
        {
            if (disc->errorf)
                (*disc->errorf)(cx, disc, ERROR_SYSTEM|2, "uniq output error");
            r = -1;
        }
        while (rsdisc(state->uniq, NiL, RS_POP));
        if (rsclose(state->uniq))
        {
            if (disc->errorf)
                (*disc->errorf)(cx, disc, ERROR_SYSTEM|2, "uniq close error");
            r = -1;
        }
        if (rskeyclose(state->uniqkey))
        {
            if (disc->errorf)
                (*disc->errorf)(cx, disc, ERROR_SYSTEM|2, "uniq key error");
            r = -1;
        }
        stkclose(state->uniqstack);
    }
    dssfclose(state->file);
    sfstrclose(state->tmp);
    vmclose(state->vm);
    return r;
}

static Cxquery_t    queries[] =
{
    {
        "sort",
        "sort records to the standard output",
        CXH,
        sort_beg,
        0,
        sort_act,
        sort_end
    },
    {0}
};

Dsslib_t        dss_lib_sort =
{
    "sort",
    "sort query"
    "[-1lms5P?\n@(#)$Id: dss sort query (AT&T Research) 2011-10-18 $\n]"
    USAGE_LICENSE,
    CXH,
    0,
    0,
    0,
    0,
    0,
    0,
    &queries[0]
};