update zebra library to be able to use links

This commit updates the zebra library files zebra.{c,h} so that it's now
possible to traverse the data structure using links! This was originally
motivated by wanting to figure out which MC particles were generated from the
MCGN bank (from which it's only possible to access the tracks and vertices
using structural links).

I've also added a new test to test-zebra which checks the consistency of all of
the next/up/orig, structural, and reference links in a zebra file.

1 files changed, 197 insertions, 45 deletions

diff --git a/src/zebra.c b/src/zebra.c
index 4ce51a2..24973d7 100644
--- a/src/zebra.c
+++ b/src/zebra.c
@@ -10,6 +10,10 @@ char zebra_err[256];
 
 zebraFile *zebra_open(const char *filename)
 {
+    /* Returns a pointer to a new zebra file object.
+     *
+     * Returns the pointer to the zebraFile object on success, or NULL on
+     * error. */
     FILE *f = fopen(filename, "r");
 
     if (!f) {
@@ -18,16 +22,23 @@ zebraFile *zebra_open(const char *filename)
     }
 
     zebraFile *z = (zebraFile *) malloc(sizeof(zebraFile));
+
+    if (!z) {
+        sprintf(zebra_err, "failed to allocate space for zebraFile struct: %s", strerror(errno));
+        fclose(f);
+        return NULL;
+    }
+
     z->f = f;
-    z->offset = 0;
     z->lr_size = 0;
     z->buf = NULL;
     z->buf_size = 0;
+    z->tab = NULL;
 
     return z;
 }
 
-int read_next_physical_record(zebraFile *z)
+static int read_next_physical_record(zebraFile *z)
 {
     /* Read the next physical record from the zebra file and append it to the
      * buffer.
@@ -94,7 +105,7 @@ int read_next_physical_record(zebraFile *z)
     return 0;
 }
 
-int get_bytes(zebraFile *z, size_t size)
+static int get_bytes(zebraFile *z, size_t size)
 {
     /* Read bytes from the zdab file until there are at least `size` words in
      * the buffer.
@@ -102,7 +113,7 @@ int get_bytes(zebraFile *z, size_t size)
      * Returns 0 on success, -1 on error, and 1 on EOF. */
     int rv;
 
-    while ((z->buf_size - z->offset) < size*4) {
+    while (z->buf_size < size*4) {
         rv = read_next_physical_record(z);
 
         if (rv == -1) {
@@ -116,65 +127,189 @@ int get_bytes(zebraFile *z, size_t size)
     return 0;
 }
 
-int read_next_logical_record(zebraFile *z)
+static uint32_t get_link(zebraFile *z, uint32_t link)
+{
+    /* Returns the offset for a zebra bank in the local buffer from a link
+     * pointing to the original bank. */
+    int i;
+    uint32_t sum;
+
+    if (z->nwtab == 0) return link;
+
+    sum = 0;
+    for (i = 0; i < z->nwtab/2; i++) {
+        if (link <= z->tab[i*2+1]) break;
+        sum += z->tab[i*2+1] - z->tab[i*2];
+    }
+
+    return sum + link - z->tab[i*2] + z->lr_offset;
+}
+
+static int rewrite_links(zebraFile *z)
+{
+    /* Rewrite the next/up/orig, structural, and reference links for all banks
+     * in the current logical record.
+     *
+     * Returns 0 on success, -1 on error. */
+    int i, rv;
+    uint32_t offset;
+    uint32_t io, noff;
+    zebraBank b;
+
+    offset = z->lr_offset*4;
+
+    while (offset < z->lr_size) {
+        io = unpacki32(z->buf+offset);
+        offset += 4;
+
+        noff = io & 0xFFFF;
+
+        if (noff > 12) {
+            offset += (noff - 12)*4;
+        }
+
+        rv = zebra_get_bank(z, &b, offset/4);
+
+        if (rv) return rv;
+
+        /* Rewrite the next, up, and orig pointers. */
+        if (b.next)
+            packi32(z->buf+offset,get_link(z,b.next));
+        if (b.up)
+            packi32(z->buf+offset+4,get_link(z,b.up));
+        if (b.orig)
+            packi32(z->buf+offset+8,get_link(z,b.orig));
+
+        /* Rewrite the structural and reference links. */
+        for (i = 0; i < b.num_links; i++) {
+            if (b.links[i])
+                packi32(z->buf+offset-(i+1)*4,get_link(z,b.links[i]));
+        }
+
+        offset += 9*4;
+
+        offset += b.num_data_words*4;
+    }
+
+    return 0;
+}
+
+int zebra_read_next_logical_record(zebraFile *z)
 {
     /* Read the next logical record into the buffer.
      *
      * Returns 0 on success, -1 on error, and 1 on EOF. */
+    int i;
     uint32_t size, type, nwtx, nwseg, nwtab, nwuh;
+    uint32_t io, noff;
+    uint32_t offset;
 
     while (1) {
-        memmove(z->buf,z->buf+z->offset,z->buf_size-z->offset);
-        z->buf = realloc(z->buf, z->buf_size-z->offset);
-        z->buf_size -= z->offset;
-        z->offset = 0;
+        /* Move the next logical record to the start of the buffer. */
+        memmove(z->buf,z->buf+z->lr_size,z->buf_size-z->lr_size);
+        z->buf = realloc(z->buf, z->buf_size-z->lr_size);
+        z->buf_size -= z->lr_size;
+        offset = 0;
 
-        switch (get_bytes(z,1)) {
+        switch (get_bytes(z,offset+1)) {
         case 1:
             return 1;
         case -1:
             return -1;
         }
-        size = unpacki32(z->buf+z->offset);
-        z->offset += 4;
-
+        size = unpacki32(z->buf+offset*4);
+        offset += 1;
+
+        /* If the size of the logical record is zero, we just skip it.
+         *
+         * From page 142 of the ZEBRA document at
+         * https://cdsweb.cern.ch/record/2296399/files/zebra.pdf:
+         *
+         * "The total size of a padding record is NWLR+1, including the LR
+         * length and the LR type; thus a padding record with NWLR=1 occupies
+         * two words. To pad exactly one word, NWLR=0 should be stored, in this
+         * case the LR type 5 is implied without being present in the data." */
         if (size == 0) continue;
 
-        switch (get_bytes(z,1)) {
+        switch (get_bytes(z,offset+1)) {
         case 1:
             return 1;
         case -1:
             return -1;
         }
-        type = unpacki32(z->buf+z->offset);
-        z->offset += 4;
+
+        /* Get the logical record type. */
+        type = unpacki32(z->buf+offset*4);
+        offset += 1;
 
         switch (type) {
+        /* Padding records. */
         case 5:
         case 6:
-            get_bytes(z,size-1);
-            z->offset += (size-1)*4;
+            get_bytes(z,offset+size-1);
+            z->lr_size = (size + 1)*4;
             continue;
+        /* Start or end of run. */
         case 1:
-            get_bytes(z,size);
-            z->offset += size*4;
+            get_bytes(z,offset+size);
+            z->lr_size = (size + 2)*4;
             continue;
+        /* Normal records. */
+        case 2:
+        case 3:
+        case 4:
+            break;
         }
 
+        /* For normal records the size of the logical record is NWLR + the two
+         * words for the size and type. */
         z->lr_size = (size + 2)*4;
-        get_bytes(z,size);
+        get_bytes(z,offset+size);
+
+        nwtx = unpacki32(z->buf+(offset+4)*4);
+        nwseg = unpacki32(z->buf+(offset+5)*4);
+        nwtab = unpacki32(z->buf+(offset+6)*4);
+        nwuh = unpacki32(z->buf+(offset+9)*4);
 
-        nwtx = unpacki32(z->buf+z->offset+4*4);
-        nwseg = unpacki32(z->buf+z->offset+5*4);
-        nwtab = unpacki32(z->buf+z->offset+6*4);
-        nwuh = unpacki32(z->buf+z->offset+9*4);
+        if (nwtab % 2) {
+            sprintf(zebra_err, "number of words in relocation table is not a multiple of 2!");
+            return -1;
+        }
 
-        z->offset += 10*4;
+        offset += 10;
 
-        z->offset += nwtx*4;
-        z->offset += nwseg*4;
-        z->offset += nwtab*4;
-        z->offset += nwuh*4;
+        /* Currently we don't use the user header, segment table, or the text
+         * vector. */
+        offset += nwuh;
+        offset += nwseg;
+        offset += nwtx;
+
+        /* Read the relocation table. */
+        if (z->tab) free(z->tab);
+        z->tab = malloc(sizeof(double)*nwtab);
+        for (i = 0; i < nwtab; i++)
+            z->tab[i] = unpacki32(z->buf+(offset+i)*4);
+        z->nwtab = nwtab;
+        offset += nwtab;
+
+        z->lr_offset = offset;
+
+        /* Find the pointer to the first bank. */
+        io = unpacki32(z->buf+offset*4);
+
+        noff = io & 0xFFFF;
+
+        /* First bank is just past the end of the logical record (one word for
+         * the io control word). */
+        z->first_bank = offset + 1;
+
+        /* Skip the links to get to the bank's next pointer. */
+        if (noff > 12) {
+            z->first_bank += noff - 12;
+        }
+
+        /* Rewrite the next/up/orig, structural, and reference links. */
+        if (rewrite_links(z)) return -1;
 
         break;
     }
@@ -182,31 +317,47 @@ int read_next_logical_record(zebraFile *z)
     return 0;
 }
 
-int next_bank(zebraFile *z, bank *b)
+int zebra_get_bank(zebraFile *z, zebraBank *b, uint32_t link)
 {
-    int rv;
-    uint32_t io, noff;
+    /* Get the bank located at `link`.
+     *
+     * On success 0 is returned and the zebraBank struct `b` is filled with the
+     * bank information. On error, -1 is returned and zebra_err is set. */
+    int i;
 
-    if (z->offset == z->lr_size) {
-        rv = read_next_logical_record(z);
-        if (rv) return rv;
+    if ((link+9)*4 > z->buf_size) {
+        sprintf(zebra_err, "link location of %i is outside buffer!", link);
+        return -1;
     }
 
-    io = unpacki32(z->buf+z->offset);
-    z->offset += 4;
+    unpack(z->buf+link*4,"lllllllll",&b->next, &b->up, &b->orig, &b->number, &b->idh, &b->num_links, &b->num_structural_links, &b->num_data_words, &b->status);
+
+    unpack(z->buf+link*4+16,"cccc",b->name,b->name+1,b->name+2,b->name+3);
 
-    noff = io & 0xFFFF;
+    b->name[4] = '\0';
 
-    if (noff > 12) {
-        z->offset += (noff-12)*4;
+    if (b->num_links > MAX_LINKS) {
+        sprintf(zebra_err, "number of links in bank is greater than %i!", MAX_LINKS);
+        return -1;
+    }
+
+    if (link < b->num_links) {
+        sprintf(zebra_err, "buffer only has %i words before bank, not enough for %i links!", link, b->num_links);
+        return -1;
     }
 
-    unpack(z->buf+z->offset,"lllllllll",&b->next, &b->up, &b->orig, &b->number, &b->name, &b->num_links, &b->num_structural_links, &b->num_data_words, &b->status);
-    z->offset += 9*4;
+    /* Read in reference and structural links. */
+    for (i = 0; i < b->num_links; i++)
+        b->links[i] = unpacki32(z->buf+(link-1-i)*4);
 
-    b->data = (uint32_t *) (z->buf+z->offset);
+    link += 9;
+
+    if ((link+b->num_data_words)*4 > z->buf_size) {
+        sprintf(zebra_err, "bank contains %i data bytes, but there are only %zu bytes left in buffer!", b->num_data_words*4, z->buf_size-link*4);
+        return -1;
+    }
 
-    z->offset += b->num_data_words*4;
+    b->data = (uint32_t *) (z->buf+link*4);
 
     return 0;
 }
@@ -214,6 +365,7 @@ int next_bank(zebraFile *z, bank *b)
 void zebra_close(zebraFile *z)
 {
     fclose(z->f);
+    if (z->tab) free(z->tab);
     if (z->buf) free(z->buf);
     free(z);
 }