Add TNG writing and reading support
[alexxy/gromacs.git] / src / external / tng_io / src / compression / huffman.c
1 /* This code is part of the tng compression routines.
2  *
3  * Written by Daniel Spangberg
4  * Copyright (c) 2010, 2013, The GROMACS development team.
5  *
6  *
7  * This program is free software; you can redistribute it and/or
8  * modify it under the terms of the Revised BSD License.
9  */
10
11
12 #include <stdio.h>
13 #include <stdlib.h>
14 #include <string.h>
15 #include "../../include/compression/warnmalloc.h"
16 #include "../../include/compression/merge_sort.h"
17 #include "../../include/compression/huffman.h"
18
19 #define MAX_HUFFMAN_LEN 31
20
21 enum htree_type { htree_leaf, htree_node };
22
23 struct htree_leaf
24 {
25   enum htree_type nodeleaf;
26   unsigned int idict; /* Index into input dictionary */
27   unsigned int prob;
28   unsigned int bit; /* One or zero */
29 };
30
31 struct htree_node
32 {
33   enum htree_type nodeleaf;
34   union htree_nodeleaf *n1;
35   union htree_nodeleaf *n2;
36   unsigned int bit; /* One or zero */
37   unsigned int prob;
38 };
39
40 union htree_nodeleaf
41 {
42   enum htree_type nodeleaf;
43   struct htree_node node;
44   struct htree_leaf leaf;
45 };
46
47 struct codelength
48 {
49   unsigned int code;
50   int length;
51   unsigned int dict;
52   unsigned int prob;
53 };
54
55 static int comp_htree(const void *leafptr1, const void *leafptr2, const void *private)
56 {
57   const union htree_nodeleaf *leaf1=(union htree_nodeleaf *)leafptr1;
58   const union htree_nodeleaf *leaf2=(union htree_nodeleaf *)leafptr2;
59   int rval=0;
60   (void)private;
61
62   if (leaf1->leaf.prob<leaf2->leaf.prob)
63     rval=1;
64   else if (leaf1->leaf.prob>leaf2->leaf.prob)
65     rval=-1;
66   return rval;
67 }
68
69 static void assign_codes(union htree_nodeleaf *htree,
70                          struct codelength *codelength,
71                          unsigned int code,
72                          int length,
73                          int top)
74 {
75 #if 0
76   printf("Assign codes called with code %d length %d\n",code,length);
77 #endif
78   if (htree->nodeleaf==htree_leaf)
79     {
80       codelength[htree->leaf.idict].length=length+1;
81       codelength[htree->leaf.idict].code=(code<<1)|htree->leaf.bit;
82 #if 0
83       printf("I am a leaf: %d %d\n",
84              codelength[htree->leaf.idict].length,
85              codelength[htree->leaf.idict].code);
86 #endif
87     }
88   else
89     {
90       if (!top)
91         {
92           code<<=1;
93           code|=htree->node.bit;
94           length++;
95         }
96 #if 0
97       printf("I am a node length: %d\n",length);
98       printf("I am a node code: %d\n",code);
99 #endif
100       assign_codes(htree->node.n1,codelength,code,length,0);
101       assign_codes(htree->node.n2,codelength,code,length,0);
102     }
103 }
104
105 static void free_nodes(union htree_nodeleaf *htree, int top)
106 {
107   if (htree->nodeleaf==htree_leaf)
108     {
109       if (!top)
110         free(htree);
111     }
112   else
113     {
114       free_nodes(htree->node.n1,0);
115       free_nodes(htree->node.n2,0);
116       if (!top)
117         free(htree);
118     }
119 }
120
121 static void flush_8bits(unsigned int *combine, unsigned char **output, int *bitptr)
122 {
123   while ((*bitptr)>=8)
124     {
125       unsigned int mask=~(0xFFU<<((*bitptr)-8));
126       unsigned char out=(unsigned char)((*combine)>>((*bitptr)-8));
127       **output=out;
128       (*output)++;
129       (*bitptr)-=8;
130       (*combine)&=mask;
131     }
132 }
133
134 static void writebits(unsigned int value,int length, unsigned char **output, int *bitptr)
135 {
136   unsigned int mask;
137   unsigned int combine=(unsigned int)**output;
138   if (length>=8)
139     mask=0xFFU<<(length-8);
140   else
141     mask=0xFFU>>(8-length);
142   while (length>8)
143     {
144       /* Make room for the bits. */
145       combine<<=8;
146       (*bitptr)+=8;
147       combine|=(value&mask)>>(length-8);
148       flush_8bits(&combine,output,bitptr);
149       length-=8;
150       mask>>=8;
151     }
152   if (length)
153     {
154       /* Make room for the bits. */
155       combine<<=length;
156       (*bitptr)+=length;
157       combine|=value;
158       flush_8bits(&combine,output,bitptr);
159     }
160   **output=(unsigned char)combine;
161 }
162
163 static unsigned int readbits(int length, unsigned char **input, int *bitptr)
164 {
165   unsigned int val=0U;
166   unsigned int extract_mask=0x80U>>*bitptr;
167   unsigned char thisval=**input;
168   while (length--)
169     {
170       val<<=1;
171       val|=((extract_mask & thisval)!=0);
172       *bitptr=(*bitptr)+1;
173       extract_mask>>=1;
174       if (!extract_mask)
175         {
176           extract_mask=0x80U;
177           *input=(*input)+1;
178           *bitptr=0;
179           thisval=**input;
180         }
181     }
182   return val;
183 }
184
185 static int comp_codes(const void *codeptr1, const void *codeptr2, const void *private)
186 {
187   const struct codelength *code1=(struct codelength *)codeptr1;
188   const struct codelength *code2=(struct codelength *)codeptr2;
189   int rval=0; /* It shouldn't be possible to get equal here, though. */
190   (void)private;
191   if (code1->length>code2->length)
192     rval=1;
193   else if (code1->length<code2->length)
194     rval=-1;
195   else if (code1->dict>code2->dict)
196     rval=1;
197   else
198     rval=-1;
199   return rval;
200 }
201
202 static int comp_codes_value(const void *codeptr1, const void *codeptr2, const void *private)
203 {
204   const struct codelength *code1=(struct codelength *)codeptr1;
205   const struct codelength *code2=(struct codelength *)codeptr2;
206
207   int rval=0; /* It shouldn't be possible to get equal here, though. */
208   (void)private;
209   if (code1->dict>code2->dict)
210     rval=1;
211   else
212     rval=-1;
213   return rval;
214 }
215
216 /* The huffman_dict array should be 131077 (0x20005) long. The
217 huffman_dict_unpacked array should be 131077 long (note five longer than
218 0x20000) */
219 void Ptngc_comp_conv_to_huffman(unsigned int *vals, int nvals,
220                           unsigned int *dict, int ndict,
221                           unsigned int *prob,
222                           unsigned char *huffman,
223                           int *huffman_len,
224                           unsigned char *huffman_dict,
225                           int *huffman_dictlen,
226                           unsigned int *huffman_dict_unpacked,
227                           int *huffman_dict_unpackedlen)
228 {
229   int i;
230   int nleft;
231   union htree_nodeleaf *htree;
232   struct codelength *codelength;
233   int bitptr;
234   unsigned char *huffman_ptr;
235   int code;
236   int longcodes=1;
237   while (longcodes)
238     {
239       /* Create array of leafs (will be array of nodes/trees during
240          buildup of tree. */
241       htree=warnmalloc(ndict*sizeof *htree);
242       codelength=warnmalloc(ndict*sizeof *codelength);
243       bitptr=0;
244       huffman_ptr=huffman;
245       for (i=0; i<ndict; i++)
246         {
247           htree[i].nodeleaf=htree_leaf;
248           htree[i].leaf.idict=i;
249           htree[i].leaf.prob=prob[i];
250         }
251       /* Sort the leafs wrt probability. */
252       Ptngc_merge_sort(htree,ndict,sizeof *htree,comp_htree,NULL);
253
254 #if 0
255       for (i=0; i<ndict; i++)
256         {
257           printf("%d %d\n",dict[htree[i].leaf.idict],htree[i].leaf.prob);
258         }
259 #endif
260
261       /* Build tree. */
262       if (ndict==1)
263         {
264           codelength[0].code=1;
265           codelength[0].length=1;
266         }
267       else
268         {
269           /* Nodes and leafs left. */
270           nleft=ndict;
271
272           /* Take the two least probable symbols (which are at the end of the
273              array and combine them until there is nothing left. */
274           while (nleft>1)
275             {
276               union htree_nodeleaf *n1=warnmalloc(sizeof *n1);
277               union htree_nodeleaf *n2=warnmalloc(sizeof *n2);
278               int new_place;
279               int p1,p2, new_prob;
280               *n1=htree[nleft-1];
281               *n2=htree[nleft-2];
282               if (n1->nodeleaf==htree_leaf)
283                 {
284                   p1=n1->leaf.prob;
285                   n1->leaf.bit=0;
286                 }
287               else
288                 {
289                   p1=n1->node.prob;
290                   n1->node.bit=0;
291                 }
292               if (n2->nodeleaf==htree_leaf)
293                 {
294                   p2=n2->leaf.prob;
295                   n2->leaf.bit=1;
296                 }
297               else
298                 {
299                   p2=n2->node.prob;
300                   n2->node.bit=1;
301                 }
302               nleft--;
303               /* Create a new node */
304               htree[nleft-1].nodeleaf=htree_node;
305               htree[nleft-1].node.n1=n1;
306               htree[nleft-1].node.n2=n2;
307               new_prob=p1+p2;
308               htree[nleft-1].node.prob=new_prob;
309               /* Use insertion sort to place this in the correct place in the
310                  array. */
311               /* Where should it be inserted? */
312               new_place=nleft;
313               while (new_place>0)
314                 {
315                   int pc;
316                   if (htree[new_place-1].nodeleaf==htree_node)
317                     pc=htree[new_place-1].node.prob;
318                   else
319                     pc=htree[new_place-1].leaf.prob;
320                   if (new_prob<pc)
321                     break;
322                   else
323                     new_place--;
324                 }
325               if (new_place!=nleft)
326                 {
327                   /* Shift array (overlapping regions!) */
328                   union htree_nodeleaf nodecopy=htree[nleft-1];
329                   memmove(htree+new_place+1,
330                           htree+new_place,
331                           (nleft-1-new_place)*sizeof *htree);
332                   htree[new_place]=nodecopy;
333                 }
334             }
335         }
336       /* Create codes from tree */
337       assign_codes(htree,codelength,0,0,1);
338       /* Canonicalize */
339       /* First put values into to the codelength array for sorting. */
340       for (i=0; i<ndict; i++)
341         {
342           codelength[i].dict=dict[i];
343           codelength[i].prob=prob[i];
344         }
345       /* Sort codes wrt length/value */
346       Ptngc_merge_sort(codelength,ndict,sizeof *codelength,comp_codes,NULL);
347       /* Canonicalize codes. */
348       code=0;
349       for (i=0; i<ndict; i++)
350         {
351           codelength[i].code=code;
352           if (i<ndict-1)
353             code=(code+1)<<(codelength[i+1].length-codelength[i].length);
354         }
355       /* Free all nodes / leaves. */
356       free_nodes(htree,1);
357       /* Free array that held nodes/leaves. */
358       free(htree);
359
360       longcodes=0;
361       /* Check if there is a too long huffman code. */
362       for (i=0; i<ndict; i++)
363         if (codelength[i].length>MAX_HUFFMAN_LEN)
364           longcodes=1;
365
366       /* If the codes are too long alter the probabilities. */
367       if (longcodes)
368         {
369           for (i=0; i<ndict; i++)
370             {
371               prob[i]>>=1;
372               if (prob[i]==0)
373                 prob[i]=1;
374             }
375
376           /* Free codelength. We will compute a new one. */
377           free(codelength);
378         }
379     }
380
381 #if 0
382   {
383     for (i=0; i<ndict; i++)
384       {
385         printf("%d %d\t\t %d %d ",codelength[i].dict,codelength[i].prob,codelength[i].length,codelength[i].code);
386         {
387           unsigned int c=codelength[i].code;
388           int j;
389           unsigned int mask=1<<(codelength[i].length-1);
390           for (j=codelength[i].length-1; j>=0; j--)
391             {
392               int bit=c&mask;
393               if (bit)
394                 printf("1");
395               else
396                 printf("0");
397               mask>>=1;
398             }
399           printf("\n");
400         }
401       }
402   }
403 #endif
404
405   /* Simply do compression by writing out the bits. */
406   for (i=0; i<nvals; i++)
407     {
408       int r;
409       for (r=0; r<ndict; r++)
410         if (codelength[r].dict==vals[i])
411           break;
412       writebits(codelength[r].code,codelength[r].length,&huffman_ptr,&bitptr);
413     }
414   if (bitptr)
415     writebits(0,8-bitptr,&huffman_ptr,&bitptr);
416   *huffman_len=(int)(huffman_ptr-huffman);
417   /* Output dictionary. */
418   /* First the largest symbol value is written in 16 bits. No bits are
419      encoded for symbols larger than this.  Then one bit signifies if
420      there is a used symbol: 1 If unused entry: 0 If used symbol the 5
421      following bits encode the length of the symbol. Worst case is
422      thus 6*65538 bits used for the dictionary. That won't happen
423      unless there's really that many values in use. If that is so,
424      well, either we compress well, or we have many values anyway. */
425   /* First sort the dictionary wrt symbol */
426   Ptngc_merge_sort(codelength,ndict,sizeof *codelength,comp_codes_value,NULL);
427   bitptr=0;
428   huffman_ptr=huffman_dict;
429   *huffman_ptr++=(unsigned char)(codelength[ndict-1].dict&0xFFU);
430   *huffman_ptr++=(unsigned char)((codelength[ndict-1].dict>>8)&0xFFU);
431   *huffman_ptr++=(unsigned char)((codelength[ndict-1].dict>>16)&0xFFU);
432   huffman_dict_unpacked[0]=(unsigned char)(codelength[ndict-1].dict&0xFFU);
433   huffman_dict_unpacked[1]=(unsigned char)((codelength[ndict-1].dict>>8)&0xFFU);
434   huffman_dict_unpacked[2]=(unsigned char)((codelength[ndict-1].dict>>16)&0xFFU);
435   for (i=0; i<=(int)codelength[ndict-1].dict; i++)
436     {
437       /* Do I have this value? */
438       int ihave=0;
439       int j;
440       for (j=0; j<ndict; j++)
441         if (codelength[j].dict==(unsigned int)i)
442           {
443
444             ihave=1;
445             writebits(1,1,&huffman_ptr,&bitptr);
446             writebits(codelength[j].length,5,&huffman_ptr,&bitptr);
447             huffman_dict_unpacked[3+i]=codelength[j].length;
448             break;
449           }
450       if (!ihave)
451         {
452           writebits(0,1,&huffman_ptr,&bitptr);
453           huffman_dict_unpacked[3+i]=0;
454         }
455     }
456   if (bitptr)
457     writebits(0,8-bitptr,&huffman_ptr,&bitptr);
458   *huffman_dictlen=(int)(huffman_ptr-huffman_dict);
459   *huffman_dict_unpackedlen=3+codelength[ndict-1].dict+1;
460
461   /* Free info about codes and length. */
462   free(codelength);
463 }
464
465 void Ptngc_comp_conv_from_huffman(unsigned char *huffman,
466                             unsigned int *vals, int nvals,
467                             int ndict,
468                             unsigned char *huffman_dict,
469                             int huffman_dictlen,
470                             unsigned int *huffman_dict_unpacked,
471                             int huffman_dict_unpackedlen)
472 {
473   struct codelength *codelength=warnmalloc(ndict*sizeof *codelength);
474   int i,j;
475   int maxdict;
476   int code;
477   unsigned char *huffman_ptr;
478   int bitptr;
479   (void)huffman_dictlen;
480   (void)huffman_dict_unpackedlen;
481   if (huffman_dict_unpacked)
482     {
483       maxdict=huffman_dict_unpacked[0]|(huffman_dict_unpacked[1]<<8)|(huffman_dict_unpacked[2]<<16);
484       j=0;
485       for(i=0; i<=maxdict; i++)
486         {
487           if (huffman_dict_unpacked[3+i]!=0)
488             {
489               codelength[j].length=huffman_dict_unpacked[3+i];
490               codelength[j].dict=i;
491 #if 0
492               printf("%d %d\n",
493                      codelength[j].length,
494                      codelength[j].dict);
495 #endif
496               j++;
497             }
498         }
499     }
500   else
501     {
502       huffman_ptr=huffman_dict;
503       maxdict=((unsigned int)huffman_ptr[0])|(((unsigned int)huffman_ptr[1])<<8)|(((unsigned int)huffman_ptr[2])<<16);
504       huffman_ptr+=3;
505       bitptr=0;
506       j=0;
507       for(i=0; i<=maxdict; i++)
508         {
509           int bit=readbits(1,&huffman_ptr,&bitptr);
510           if (bit)
511             {
512               codelength[j].length=readbits(5,&huffman_ptr,&bitptr);
513               codelength[j].dict=i;
514 #if 0
515               printf("%d %d\n",
516                      codelength[j].length,
517                      codelength[j].dict);
518 #endif
519               j++;
520             }
521         }
522     }
523   /* Sort codes wrt length/value. */
524   Ptngc_merge_sort(codelength,ndict,sizeof *codelength,comp_codes,NULL);
525   /* Canonicalize codes. */
526   code=0;
527   for (i=0; i<ndict; i++)
528     {
529       codelength[i].code=code;
530       if (i<ndict-1)
531         code=(code+1)<<(codelength[i+1].length-codelength[i].length);
532     }
533 #if 0
534   {
535     for (i=0; i<ndict; i++)
536       {
537         printf("%d -\t\t %d %d ",codelength[i].dict,codelength[i].length,codelength[i].code);
538         {
539           unsigned int c=codelength[i].code;
540           int j;
541           unsigned int mask=1<<(codelength[i].length-1);
542           for (j=codelength[i].length-1; j>=0; j--)
543             {
544               int bit=c&mask;
545               if (bit)
546                 printf("1");
547               else
548                 printf("0");
549               mask>>=1;
550             }
551           printf("\n");
552         }
553       }
554   }
555 #endif
556   /* Decompress data. */
557   huffman_ptr=huffman;
558   bitptr=0;
559   for (i=0; i<nvals; i++)
560     {
561       unsigned int symbol;
562       int len=codelength[0].length;
563       symbol=readbits(len,&huffman_ptr,&bitptr);
564       j=0;
565       while (symbol!=codelength[j].code)
566         {
567           int newlen;
568           j++;
569           newlen=codelength[j].length;
570           if (newlen!=len)
571             {
572               symbol<<=(newlen-len);
573               symbol|=readbits(newlen-len,&huffman_ptr,&bitptr);
574               len=newlen;
575             }
576         }
577       vals[i]=codelength[j].dict;
578     }
579   /* Free info about codes and length. */
580   free(codelength);
581 }