1
\$\begingroup\$

Some time ago, I implemented shuf(1), but it had an issue: it ignored nul characters in the input.

This new version tries to solve this problem (but uses the double of memory).

shuf(1) shuffles the lines of the files given as arguments or the standard input (if no argument was provided) and print it. If the -n NUM option were given, prints NUM random lines from the files.

#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <stdint.h>
#include <string.h>
#include <unistd.h>

#ifdef LIBBSD
#include <bsd/stdlib.h>
#include <bsd/unistd.h>
#endif

/* The buffer structure */
struct Buffer {
    char  *chars;       /* array of characters in the buffer */
    char **linebeg;     /* array of pointers to beginning of lines in chars[] */
    char **lineend;     /* array of pointers to end of lines in chars[] */
    size_t maxchars;    /* size of the array of characters */
    size_t maxlines;    /* size of the array of lines */
    size_t nchars;      /* number of characters in the buffer */
    size_t nlines;      /* number of lines in the buffer */
};

static int readfile(int fd, struct Buffer *buf);
static void getlines(struct Buffer *buf);
static void shuf(struct Buffer *buf, size_t nshuf);
static void printbuf(struct Buffer *buf, size_t nshuf);
static size_t getnshuf(const char *s);
static void usage(void);

/* sortnsearch: read a file into memory, sort it by lines and search a line */
int
main(int argc, char *argv[])
{
    int c;
    int fd;
    int exitval;
    size_t nshuf;
    char *filename;
    struct Buffer buf = {NULL, NULL, NULL, 0, 0, 0, 0};

    nshuf = 0;
    while ((c = getopt(argc, argv, "n:")) != -1) {
        switch (c) {
        case 'n':
            nshuf = getnshuf(optarg);
            if (nshuf == 0)
                errx(EXIT_FAILURE, "%s: invalid number of lines", optarg);
            break;
        default:
            usage();
            break;
        }
    }
    argc -= optind;
    argv += optind;

    exitval = EXIT_SUCCESS;

    fd = STDIN_FILENO;
    filename = "stdin";
    do {
        if (*argv) {
            if (**argv == '-' && **(argv+1) == '\0') {
                fd = STDIN_FILENO;
            } else if ((fd = open(*argv, O_RDONLY)) == -1) {
                warn("%s", *argv);
                exitval = EXIT_FAILURE;
                argv++;
                continue;
            }
            filename = *argv++;
        }
        if (readfile(fd, &buf) == -1) {
            warn("%s", *argv);
            exitval = EXIT_FAILURE;
        }
        if (fd != STDIN_FILENO)
            (void)close(fd);
    } while (*argv);

    getlines(&buf);

    if (nshuf == 0 || nshuf > buf.nlines)
        nshuf = buf.nlines;

    shuf(&buf, nshuf);

    printbuf(&buf, nshuf);

    free(buf.chars);
    free(buf.linebeg);
    free(buf.lineend);

    return exitval;
}

/* read file fd into memory in buf->chars */
static int
readfile(int fd, struct Buffer *buf)
{
    char tmpbuf[BUFSIZ], *tmp;
    size_t size = buf->maxchars;
    ssize_t n;

    while ((n = read(fd, tmpbuf, sizeof tmpbuf)) != 0 && n != -1) {
        if (n + buf->nchars >= size) {
            size = buf->maxchars + sizeof tmpbuf;

            /* overflow check */
            if (size <= buf->maxchars) {
                (void)free(buf->chars);
                errc(EXIT_FAILURE, EOVERFLOW, NULL);
            }

            if ((tmp = realloc(buf->chars, size)) == NULL) {
                (void)free(buf->chars);
                err(EXIT_FAILURE, "realloc");
            }
            buf->chars = tmp;
            buf->maxchars = size;
        }
        memcpy(buf->chars + buf->nchars, tmpbuf, n);
        buf->nchars += n;
    }

    if (n == -1)
        return -1;

    return 1;
}

/* get beginnings end ends of lines from buf->chars into buf->linebeg and buf->lineend */
static void
getlines(struct Buffer *buf)
{
    size_t i;
    size_t size;
    char **tmp;

    if ((buf->linebeg = reallocarray(NULL, BUFSIZ, sizeof *buf->linebeg)) == NULL)
        err(EXIT_FAILURE, "realloc");
    if ((buf->lineend = reallocarray(NULL, BUFSIZ, sizeof *buf->lineend)) == NULL)
        err(EXIT_FAILURE, "realloc");
    buf->linebeg[buf->nlines++] = buf->chars;
    buf->maxlines = size = BUFSIZ;

    for (i = 0; i < buf->nchars; i++) {
        if (buf->chars[i] == '\n' || i == buf->nchars - 1) {
            if (buf->nlines + 1 >= size) {
                size = buf->maxlines << 1;

                if ((tmp = reallocarray(buf->linebeg, size, sizeof *buf->linebeg)) == NULL)
                    err(EXIT_FAILURE, "realloc");
                buf->linebeg = tmp;
                if ((tmp = reallocarray(buf->lineend, size, sizeof *buf->lineend)) == NULL)
                    err(EXIT_FAILURE, "realloc");
                buf->lineend = tmp;
                buf->maxlines = size;
            }
            buf->lineend[buf->nlines-1] = buf->chars+i+1;
            buf->linebeg[buf->nlines++] = buf->chars+i+1;
        }
    }
    buf->lineend[--buf->nlines] = buf->chars+i+1;
}

/* shuffle buf->linebeg and buf->lineend parallelly */
static void
shuf(struct Buffer *buf, size_t nshuf)
{
    size_t i, j;
    char *tmp;

    for (i = buf->nlines-1; i > 0 && nshuf > 0; i--, nshuf--) {
        j = arc4random_uniform(i + 1);

        tmp = buf->linebeg[i];
        buf->linebeg[i] = buf->linebeg[j];
        buf->linebeg[j] = tmp;

        tmp = buf->lineend[i];
        buf->lineend[i] = buf->lineend[j];
        buf->lineend[j] = tmp;
    }
}

/* print at most nshuf lines */
static void
printbuf(struct Buffer *buf, size_t nshuf)
{
    while (nshuf-- > 0)
        write(STDOUT_FILENO, buf->linebeg[nshuf],
              buf->lineend[nshuf] - buf->linebeg[nshuf]);
}

/* return number from s */
static size_t
getnshuf(const char *s)
{
    size_t n;
    char *endp;

    n = strtoull(s, &endp, 10);
    if (errno == ERANGE || endp == s || *endp != '\0')
        return 0;

    return (size_t) n;
}

/* show usage */
static void
usage(void)
{
    (void)fprintf(stderr, "usage: shuf [-n number] [file...]\n");
    exit(EXIT_FAILURE);
}
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.