binoculars/hkl-binoculars-cnpy.c

   1 #include <assert.h>
   2 #include <stdint.h>
   3 #include <stdio.h>
   4 #include <stdlib.h>
   5 #include <sys/types.h>
   6 #include <regex.h>
   7 #include <string.h>
   8
   9 #include "hkl-binoculars-cnpy-private.h"
  10 #include "hkl/ccan/array_size/array_size.h"
  11 #include "hkl/ccan/build_assert/build_assert.h"
  12 #include "hkl/ccan/darray/darray.h"
  13
  14 #define REGEX_DESCR "'descr':\\s+'(.+)'"
  15 #define REGEX_FORTRAN "'fortran_order':\\s+(.+)"
  16 #define REGEX_SHAPE "'shape':\\s+\\((.+)\\)"
  17 #define REGEX_HEADER "^\\{" REGEX_DESCR ",\\s+" REGEX_FORTRAN ",\\s+" REGEX_SHAPE ",\\s+" "\\}\\s+"
  18
  19 datatype(
  20         Endianess,
  21         (LittleEndian),
  22         (BigEndian)
  23         );
  24
  25 static const char magic[] = {(char)0x93, 'N', 'U', 'M', 'P', 'Y'};
  26
  27 const uint8_t VERSION_1 = 0x01;
  28 const uint8_t VERSION_2 = 0x02;
  29 const uint8_t VERSION_3 = 0x03;
  30
  31 struct pre_header_t {
  32         char magic[6];
  33         uint8_t major;
  34         uint8_t minor;
  35 };
  36
  37 struct descr_t {
  38         Endianess endianess;
  39         HklBinocularsNpyDataType elem_type;
  40         size_t elem_size;
  41 };
  42
  43 struct npy_t {
  44         struct pre_header_t pre;
  45         uint32_t header_len;
  46         char *header;
  47         struct descr_t descr;
  48         int fortran_order;
  49         darray_int shape;
  50         void *arr;
  51 };
  52
  53 void npy_free_but_array(struct npy_t *npy)
  54 {
  55         free(npy->header);
  56         free(npy);
  57 }
  58
  59 static char *extract_as_string(const char *header, regmatch_t match)
  60 {
  61         int size = match.rm_eo - match.rm_so;
  62
  63         char *buffer = g_new(char, size + 1);
  64         strncpy(buffer, &header[match.rm_so], size);
  65         buffer[size] = 0x0;
  66
  67         return buffer;
  68 }
  69
  70 static struct descr_t parse_descr(const char *header, regmatch_t match)
  71 {
  72         char *description = extract_as_string(header, match);
  73
  74         struct descr_t descr;
  75
  76         /* endianness */
  77         switch(description[0]){
  78         case '|':
  79         case '<':
  80                 descr.endianess = LittleEndian();
  81                 break;
  82         };
  83
  84         /* elem_type */
  85         switch(description[1]){
  86         case 'b':
  87                 descr.elem_type = HklBinocularsNpyBool();
  88                 break;
  89         };
  90
  91         /* elem_size */
  92         descr.elem_size = atoi(&description[2]);
  93
  94         free(description);
  95
  96         return descr;
  97 }
  98
  99 static int parse_fortran_order(const char *header, regmatch_t match)
 100 {
 101         int fortran_order = 0;
 102
 103         if(!strncmp(&header[match.rm_so], "True", 4))
 104                 fortran_order = 1;
 105
 106         return fortran_order;
 107 }
 108
 109
 110 static void parse_shape(const char *header, regmatch_t match, darray_int *shape)
 111 {
 112         char *repr = extract_as_string(header, match);
 113         char *token;
 114         char *saveptr;
 115
 116         darray_free(*shape);
 117         darray_init(*shape);
 118
 119         for(;;repr=NULL){
 120                 int val;
 121
 122                 token=strtok_r(repr, ",", &saveptr);
 123                 if(NULL == token)
 124                         break;
 125
 126                 val = atoi(token);
 127                 /* TODO  deal with atoi errors */
 128                 darray_append(*shape, val);
 129         };
 130
 131         free(repr);
 132 }
 133
 134 static int shape_cmp(const darray_int *sh1, const darray_int *sh2)
 135 {
 136         int i;
 137
 138         if (darray_size(*sh1) != darray_size(*sh2))
 139                 return -1; /* not idential */
 140
 141         for(i=0; i<darray_size(*sh1); ++i)
 142                 if(darray_item(*sh1, i) != darray_item(*sh2, i))
 143                         return -1;
 144
 145         return 0; /* idential */
 146 }
 147
 148 static int shape_size(const darray_int *shape)
 149 {
 150         int nb_elem = 1;
 151         int *val;
 152
 153         darray_foreach(val, *shape){
 154                 nb_elem *= *val;
 155         }
 156
 157         return nb_elem;
 158 }
 159
 160 static struct npy_t *parse_npy(FILE* fp,
 161                                HklBinocularsNpyDataType type,
 162                                const darray_int *shape)
 163 {
 164         struct npy_t *npy = g_new0(struct npy_t, 1);
 165         size_t res;
 166         regex_t preg;
 167         regmatch_t matches[4];
 168         int errcode;
 169         char errbuf[256];
 170
 171         BUILD_ASSERT(sizeof(npy->pre) == 8);
 172
 173         /* read the pre_header */
 174         res = fread(&npy->pre, 1, sizeof(npy->pre), fp);
 175         assert(res == 8);
 176
 177         if(strncmp(magic, npy->pre.magic, ARRAY_SIZE(magic)))
 178                 goto fail_no_header;
 179
 180         /* read the header size */
 181         switch (npy->pre.major) {
 182         case VERSION_1:
 183         {
 184                 uint16_t len;
 185                 res = fread(&len, 1, sizeof(len), fp);
 186                 assert(res == 2);
 187                 npy->header_len = (uint32_t)len;
 188                 break;
 189         }
 190         case VERSION_2:
 191         {
 192                 res = fread(&npy->header_len, 1, sizeof(npy->header_len), fp);
 193                 assert(res == 4);
 194                 break;
 195         }
 196         case VERSION_3:
 197         {
 198                 res = fread(&npy->header_len, 1, sizeof(npy->header_len), fp);
 199                 assert(res == 4);
 200                 break;
 201         }
 202         };
 203
 204         /* read the header */
 205
 206         npy->header = malloc(npy->header_len * sizeof(char));
 207         res = fread(npy->header, 1, npy->header_len, fp);
 208         assert(res == npy->header_len);
 209
 210
 211         /* parse the header */
 212         /* {descr: '|b1', fortran_order: False, shape: (240, 560), } */
 213
 214         errcode = regcomp(&preg, REGEX_HEADER, REG_EXTENDED);
 215         if(errcode){
 216                 regerror(errcode, &preg, errbuf, 256);
 217                 fprintf(stdout, "errcode: %s\n", errbuf);
 218
 219                 goto fail;
 220         }
 221
 222         errcode = regexec(&preg, npy->header, ARRAY_SIZE(matches), matches, 0);
 223         if(errcode){
 224                 regerror(errcode, &preg, errbuf, 256);
 225                 fprintf(stdout, "errcode: %s\n", errbuf);
 226
 227                 regfree(&preg);
 228                 goto fail;
 229         }
 230
 231         regfree(&preg);
 232
 233         npy->descr = parse_descr(npy->header, matches[1]);
 234         npy->fortran_order = parse_fortran_order(npy->header, matches[2]);
 235         parse_shape(npy->header, matches[3], &npy->shape);
 236
 237         if(type.tag != npy->descr.elem_type.tag)
 238                 goto fail;
 239
 240         if(0 != shape_cmp(shape, &npy->shape))
 241                 goto fail;
 242
 243         /* read the array */
 244         int nbytes = shape_size(&npy->shape) * npy->descr.elem_size;
 245         npy->arr = malloc( nbytes );
 246         if(NULL == npy->arr)
 247                 goto fail;
 248
 249         res = fread(npy->arr, 1, nbytes, fp);
 250         assert(res == nbytes);
 251
 252         return npy;
 253 fail:
 254         free(npy->header);
 255 fail_no_header:
 256         free(npy);
 257
 258         return NULL;
 259 }
 260
 261 void *npy_load(const char *fname,
 262                HklBinocularsNpyDataType type,
 263                const darray_int *shape)
 264 {
 265         uint8_t *arr = NULL;
 266         FILE* fp = fopen(fname, "rb");
 267
 268         if (NULL != fp){
 269                 struct npy_t *npy = parse_npy(fp, type, shape);
 270                 if (NULL != npy) {
 271                         arr = npy->arr;
 272                         npy_free_but_array(npy);
 273                 }
 274
 275                 fclose(fp);
 276         }
 277         return arr;
 278 }
 279
 280 static inline char bigendian(void)
 281 {
 282         int x = 1;
 283         return (((char *)&x)[0]) ? '<' : '>';
 284 }
 285
 286 static inline char map_type(HklBinocularsNpyDataType type)
 287 {
 288         char res = '?';
 289
 290         match(type){
 291                 of(HklBinocularsNpyBool)   res = 'b';
 292                 of(HklBinocularsNpyDouble) res = 'f';
 293         }
 294
 295         /* if(t == typeid(float) ) return 'f'; */
 296         /* if(t == typeid(double) ) return 'f'; */
 297         /* if(t == typeid(long double) ) return 'f'; */
 298
 299         /* if(t == typeid(int) ) return 'i'; */
 300         /* if(t == typeid(char) ) return 'i'; */
 301         /* if(t == typeid(short) ) return 'i'; */
 302         /* if(t == typeid(long) ) return 'i'; */
 303         /* if(t == typeid(long long) ) return 'i'; */
 304
 305         /* if(t == typeid(unsigned char) ) return 'u'; */
 306         /* if(t == typeid(unsigned short) ) return 'u'; */
 307         /* if(t == typeid(unsigned long) ) return 'u'; */
 308         /* if(t == typeid(unsigned long long) ) return 'u'; */
 309         /* if(t == typeid(unsigned int) ) return 'u'; */
 310
 311         /* if(t == typeid(bool) ) return 'b'; */
 312
 313         /* if(t == typeid(std::complex<float>) ) return 'c'; */
 314         /* if(t == typeid(std::complex<double>) ) return 'c'; */
 315         /* if(t == typeid(std::complex<long double>) ) return 'c'; */
 316
 317         /* else return '?'; */
 318         return res;
 319 }
 320
 321 static inline int map_size(HklBinocularsNpyDataType type)
 322 {
 323         int res = 8;
 324
 325         match(type){
 326                 of(HklBinocularsNpyBool)   res = 1;
 327                 of(HklBinocularsNpyDouble) res = 8;
 328         }
 329
 330         return res;
 331 }
 332
 333 static inline char *create_npy_header(HklBinocularsNpyDataType type,
 334                                       const darray_int *shape,
 335                                       size_t *header_size)
 336 {
 337         int i;
 338         char *header;
 339         char *dict;
 340         size_t dict_size;
 341         FILE *stream;
 342
 343         /* first compute the dict and its size */
 344
 345         stream = open_memstream(&dict, &dict_size);
 346         fprintf(stream,"{");
 347         fprintf(stream, "'descr': '%c%c%d'",
 348                 bigendian(), map_type(type), map_size(type));
 349         fprintf(stream, ", ");
 350         fprintf(stream, "'fortran_order' : False");
 351         fprintf(stream, ", ");
 352         fprintf(stream, "'shape': (");
 353         fprintf(stream, "%d", darray_item(*shape, 0));
 354         for(i=1; i<darray_size(*shape); ++i){
 355                 fprintf(stream, ", %d", darray_item(*shape, i));
 356         }
 357         fprintf(stream, ")");
 358         fprintf(stream, ",");
 359         fprintf(stream, "}");
 360
 361         fflush(stream);
 362
 363         /* ``len(magic string) + 2 + len(length) + HEADER_LEN`` be
 364            evenly divisible by 64 for alignment purposes. */
 365         size_t extra = 64 - (ARRAY_SIZE(magic) + 2 + 2 + dict_size) % 64;
 366         for(i=0; i<extra - 1; ++i){
 367                 fprintf(stream, " ");
 368         }
 369         fprintf(stream, "\n");
 370
 371         fclose(stream);
 372
 373         /* create the full header */
 374
 375         stream = open_memstream(&header, header_size);
 376         fwrite(&magic[0], sizeof(char), ARRAY_SIZE(magic), stream);
 377         fprintf(stream, "%c%c", VERSION_1, 0x0);
 378         uint16_t len = (uint16_t)dict_size;
 379         fwrite(&len, 1, sizeof(len), stream);
 380         fwrite(dict, 1, dict_size, stream);
 381         fprintf(stdout, "%s", dict);
 382         fflush(stream);
 383         fclose(stream);
 384
 385         free(dict);
 386
 387         return header;
 388 }
 389
 390 void npy_save(const char *fname,
 391               const void *arr,
 392               HklBinocularsNpyDataType type,
 393               const darray_int *shape)
 394 {
 395         FILE *f;
 396         char *header;
 397         size_t header_size;
 398
 399         header = create_npy_header(type, shape, &header_size);
 400         f = fopen(fname, "wb");
 401         fseek(f,0,SEEK_SET);
 402         fwrite(&header[0], 1, header_size, f);
 403         fprintf(stdout, "%s\n", header);
 404         fseek(f,0,SEEK_END);
 405         fwrite(arr, map_size(type), shape_size(shape), f);
 406         fclose(f);
 407
 408         free(header);
 409 }