; _______________________________________________________________________________________
;|                                                                                       |
;| ..::FreshLib::..  Free, open source. Licensed under "BSD 2-clause" license."          |
;|_______________________________________________________________________________________|
;
;  Description: Deflate uncompress library.
;
;  Target OS: Any
;
;  Dependencies:
;
;  Notes:
;_________________________________________________________________________________________

module "Deflate"


DEFLATE_MAX_BITS = 15

DEFLATE_SIZE_LEN_TREE  = 320*2
DEFLATE_SIZE_DIST_TREE = 32*2
DEFLATE_SIZE_LL_TREE   = 32*2
DEFLATE_SIZE_LEN_BUF   = 300+300


; Decompress the block in memory.
; Arguments:
;  .pDest    points to a buffer of type TBitStream for the uncompressed data.
;  .pSrc     points to a source memory buffer.
;  .bitOffs  bit offset of the start of the stream, related to [.pSrc].
;  .srcCount byte size of the source buffer, measured from [.pSrc].
;
; Returns:
;  CF=0; The stream is decompressed without errors.
;
;  CF=1: There was error on decompression, because of malformed stream or memory
;        allocations.
;
;  EAX:  Pointer to the destination stream. (It can be changed on reallocations)
;
;  NOTE: Even on error, the stream can be partially decompressed in the destination
;        stream;
;
proc Inflate, .pDest, .dstCount, .pSrc, .srcCount

  .len_tree      rd DEFLATE_SIZE_LEN_TREE
  .dist_tree     rd DEFLATE_SIZE_DIST_TREE
  .ll_tree       rd DEFLATE_SIZE_LL_TREE
  .lengths       rb DEFLATE_SIZE_LEN_BUF

  .final         dd ?

  .pEndSrc       dd ?
  .pEndDst       dd ?
  .pCurrent      dd ?

  .p_lengths_end dd ?

  .LLCount       dd ?
  .DCount        dd ?

  .b_type dd ?

begin
        pushad

        mov     eax, [.pDest]
        mov     [.pCurrent], eax
        add     eax, [.dstCount]
        mov     [.pEndDst], eax

; prepare the stream pointers.

        mov     esi, [.pSrc]

        mov     eax, esi
        add     eax, [.srcCount]
        mov     [.pEndSrc], eax

        mov     ch, $10
        mov     eax, [esi]

.main_loop:

; read 1 bit header BFINAL

        mov     [.final], eax
        shr     eax, 1

        and     [.final], 1

        dec     ch
        jnz     @f

        add     esi, 2
        cmp     esi, [.pEndSrc]
        jae     .error2

        mov     eax, [esi]
        mov     ch, $10

@@:

; read 2 bit header BTYPE
        mov     edx, eax
        shr     eax, 2
        and     edx, 3

        sub     ch, 2

        dec     edx
        js      .uncompressed     ; BTYPE == 0 uncompressed type

        mov     [.b_type], edx    ; 0 if fixed tree; <>0 if dynamic tree
        jnz     .dynamic_tree

; BTYPE == 1 compressed with fixed huffman tree.

; load the fixed huffman tree for literal/lengths.
        pushad

        lea     edi, [.lengths]
        lea     esi, [.TreeFixed]
        mov     ebx, edi

        xor     eax, eax
        xor     ecx, ecx
        mov     dl, 4      ; edx == 0 here.

.loop_fx:
        lodsw
        mov     cl, ah
        rep stosb
        dec     edx
        jnz     .loop_fx

        sub     edi, ebx
        lea     eax, [.len_tree]

        stdcall __tree_by_lengths, ebx, edi, eax

        popad

        jmp     .decode


.dynamic_tree:
        dec     edx
        jnz     .error3           ; BTYPE == 3 invalid block type

; Compressed with dynamic huffman table. Read it from the stream.

; first read the counts

; 5 bits HLIT - Number of Literal/length codes - 257.
; 5 bits HDIST - Number of distance codes - 1.
; 4 bits HCLEN - Number of code length codes - 4.

        mov     edi, eax
        mov     ebx, eax
        mov     edx, eax

        shr     ebx, 5
        shr     edx, 10
        shr     eax, 14

        and     edi, $1f
        and     ebx, $1f
        and     edx, $0f

        sub     ch, 14
        jg      @f

        add     esi, 2
        cmp     esi, [.pEndSrc]
        jae     .error4

        mov     cl, ch
        mov     eax, [esi]
        add     ch, $10
        rol     eax, cl

@@:
        add     edi, 257
        inc     ebx
        add     edx, 4

        mov     [.LLCount], edi
        mov     [.DCount], ebx

; second, read the fixed 3bit table for the code lengths.

; clear the lengths array

        lea     edi, [.lengths]
        xor     ebx, ebx

        mov     [edi],    ebx
        mov     [edi+4],  ebx
        mov     [edi+8],  ebx
        mov     [edi+12], ebx
;        mov     [edi+16], ebx  ; codes 16, 17 and 18 are always transmitted, no need to clear;
        mov     bl, $0f

.loop_hc:
        movzx   ebx, [.clclOrder2+ebx]

        mov     [edi+ebx], al
        shr     eax, 3

        and     byte [edi+ebx], $07

        sub     ch, 3
        jg      @f

        add     esi, 2
        cmp     esi, [.pEndSrc]
        jae     .error5

        mov     cl, ch
        mov     eax, [esi]
        add     ch, $10
        rol     eax, cl

@@:
        dec     edx
        jnz     .loop_hc

        lea     ebx, [.ll_tree]
        stdcall __tree_by_lengths, edi, 19, ebx
        jnz     .error6

; now create the two remaining huffman trees: for literal/lengths and for distances.

; First the literal/lengths huffman tree

        mov     edx, [.LLCount]
        lea     edi, [.len_tree]
        call    .decode_lengths
        jnz     .error9

; Second, the distances huffman tree.

        mov     edx, [.DCount]
        lea     edi, [.dist_tree]
        call    .decode_lengths
        jnz     .error10

; At the end - trees are ready, so decode the compressed block.

        mov     edi, [.pCurrent]

.decode:
        cmp     esi, [.pEndSrc]
        jae     .error7

; read literal/length code

        xor     edx, edx

.bitloop1:
        shr     eax, 1
        rcl     edx, 1                    ; edx = 2*edx+CF

        dec     ch
        jnz     @f

        add     esi, 2

        mov     eax, [esi]
        mov     ch, $10

@@:
        xor     edx, [.len_tree+4*edx]
        jns     .bitloop1

        not     edx

        cmp     edx, 255
        ja      .copy_from_older        ; if EDX <= 255 then it is simple encoded byte, so store it.


;        OutputValue "Literal: ", edx, 16, 2


; Store literal value.
        cmp     edi, [.pEndDst]
        jae     .error14

        mov     [edi], dl
        inc     edi
        jmp     .decode



.copy_from_older:
;        OutputValue "Length code:", edx, 10, 3

        sub     edx, 257
        js      .next_block             ; if EDX == 256 then end of block flag.

; The code is a length code. Compute the length of string to be copied from the previous decoded data.

; Compute the additional length and the extra bit count.
; 2^n(4 + edx mod 4)

        xor     ebx, ebx

        cmp     edx, 4
        jb      .length_ok

        sub     edx, 4
        mov     cl, dl
        and     edx, 3  ; edx = edx mod 4
        shr     cl, 2   ; cl = edx div 4
        lea     edx, [edx+4]
        jz      .length_ok
        shl     edx, cl

        dec     edx
        cmp     cl, 6           ; damn exceptions!
        jae     .length_ok

        adc     ebx, edx        ; CF = 1 here!

; read CL bits from the stream
        or      edx, -1
        shl     edx, cl
        not     edx

        and     edx, eax
        shr     eax, cl

;        OutputValue "Additional value:", edx, 10, 2

        sub     ch, cl
        jg      .length_ok

        add     esi, 2

        mov     cl, ch
        mov     eax, [esi]
        add     ch, $10
        rol     eax, cl

.length_ok:
        lea     ebx, [ebx+edx+3]


;        OutputValue "Copy length:", ebx, 10, 3

; ... so now decode the distance code

        cmp     [.b_type], 0
        jne     .decode_dist  ; when there is fixed huffman table, then the distances are encoded with 5 bit fixed code

; read fixed 5 bit distance.

        mov     cl, 5
        xor     edx, edx

.bitloop4:
        shr     eax, 1
        rcl     edx, 1

        dec     cl
        jnz     .bitloop4

;        OutputValue "Distance code:", edx, 10, 2

        sub     ch, 5
        jg      .compute_distance

        add     esi, 2

        mov     cl, ch
        mov     eax, [esi]
        add     ch, $10
        rol     eax, cl

        jmp     .compute_distance


.decode_dist:
        xor     edx, edx

.bitloop2:
        shr     eax, 1
        rcl     edx, 1                    ; edx = 2*edx+CF

        dec     ch
        jnz     @f

        add     esi, 2

        mov     eax, [esi]
        mov     ch, $10

@@:
        xor     edx, [.dist_tree+4*edx]
        jns     .bitloop2

        not     edx

.compute_distance:
        cmp     edx, 30
        jae     .error8

; 2^n(2 + edx mod 2)
; compute the additional distance and extra bit count for the distance.

        xor     cl, cl
        cmp     edx, 2
        jb      .distance_ok

        sub     edx, 2
        mov     cl, dl  ;
        and     edx, 1  ; edx = edx mod 2
        shr     cl, 1   ; cl = edx div 2
        lea     edx, [edx+2]
        jz      .distance_ok

        shl     edx, cl  ; edx = 2^cl * (edx mod 2 + 2)

; read CL bits from the stream in EBX
        push    ebx

        or      ebx, -1  ; mask
        shl     ebx, cl  ; mask
        not     ebx      ; mask

        and     ebx, eax
        shr     eax, cl

        add     edx, ebx
        pop     ebx

        sub     ch, cl
        jg      .distance_ok

        add     esi, 2

        mov     cl, ch
        mov     eax, [esi]
        add     ch, $10
        rol     eax, cl

.distance_ok:
        inc     edx

; copy the string
        push    esi ecx

        neg     edx
        mov     ecx, ebx

        lea     esi, [edi+edx]
        lea     ebx, [edi+ecx]

        cmp     esi, [.pDest]
        jb      .error_cp1

        cmp     ebx, [.pEndDst]
        ja      .error_cp2

        rep movsb

        pop     ecx esi
        jmp     .decode

.error_cp1:
        dbrk
.error_cp2:
        dbrk
        pop     ecx esi

; these all "error" labels are equal in release mode. In debug mode, there are int3 instructions inbetween in order to allow error tracing.
.error2:
        dbrk
.error3:
        dbrk
.error4:
        dbrk
.error5:
        dbrk
.error6:
        dbrk
.error7:
        dbrk
.error8:
        dbrk
.error9:
        dbrk
.error10:
        dbrk
.error11:
        dbrk
.error12:
        dbrk
.error13:
        dbrk
.error14:
        dbrk
.error:
        dbrk
        not     [.final]          ; Tricky makes CF=1 on exit.

.next_block:
        mov     [.pCurrent], edi
        cmp     [.final], 0
        je      .main_loop

        shl     [.final], 1      ; If normal end [.final] is always 1; If error, [.final] is $ffffffff or $fffffffe

        popad
        return


; not compressed block
;-----------------------------------------------------------
.uncompressed:

; align to byte
        movzx   ecx, ch
        shr     ecx, 3
        add     esi, ecx

        cmp     esi, [.pEndSrc]
        jae     .error11

        mov     edi, [.pCurrent]

        xor     eax, eax

        lodsw                           ; LEN
        mov     ecx, eax
        lodsw                           ; NLEN
        not     ax
        cmp     ecx, eax
        jne     .error12

        add     eax, edi
        cmp     eax, [.pEndDst]
        ja      .error13

        rep movsb

; resync stream
        mov     ch, $10
        mov     eax, [esi]
        jmp     .next_block


; ebx = pointer to the .ll_tree huffman tree
; edx = count of the encoded lengths
; edi = pointer to the tree memory

.decode_lengths:

; arguments.
        push    edi
        push    edx
        lea     edi, [.lengths]
        push    edi

; buffer overflow?
        add     edx, edi
        mov     [.p_lengths_end], edx

.loop_len:
        xor     edx, edx
        cmp     esi, [.pEndSrc]
        jae     .errordl

.bitloop3:
        shr     eax, 1
        rcl     edx, 1                  ; edx = 2*edx+CF

        dec     ch
        jnz     @f

        add     esi, 2
        mov     eax, [esi]
        mov     ch, $10

@@:
        xor     edx, [ebx+4*edx]
        jns     .bitloop3

; here dh = $ff
        xor     dh, dl          ; dh = not dl because dh is always $ff for this tree.

        mov     dl, 1           ; count
        cmp     dh, 16
        jb      .rep_loop

        mov     cl, 7
        mov     dl, 11
        sub     dh, 18
        jz      .repeat         ; dh = 18

        mov     cl, 3
        mov     dl, 3
        inc     dh
        jz      .repeat         ; dh = 17

        cmp     edi, [esp]
        je      .errordl        ; it is the first length, there is no previous.

        mov     cl, 2
        mov     dh, [edi-1]   ; previous length; dl=3 here

.repeat:
        shl     edx, 8

; read CL bits from the stream: CL = 2, 3 or 7
        or      dl, -1
        shl     dl, cl
        not     dl

        and     dl, al
        shr     eax, cl

        sub     ch, cl
        jg      @f

        add     esi, 2

        mov     cl, ch
        mov     eax, [esi]
        add     ch, $10
        rol     eax, cl

@@:
        add     dh, dl
        shr     edx, 8          ; now dl is the count, dh is the byte to be repeat!

.rep_loop:
        mov     [edi], dh
        inc     edi
        dec     dl
        jnz     .rep_loop

; next length
        cmp     edi, [.p_lengths_end]
        jb      .loop_len
        ja      .errordl

        stdcall __tree_by_lengths   ; the arguments are already in the stack.
        retn

; The assembly way of error handling. ;)
.errordl:
        add     esp, 3*4 + 4 ; pushed arguments and the return address.
        jmp     .error

;  .clclOrder     db 16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15   ; original order.
.clclOrder2  db  8, 15, 14, 13, 12, 11, 10, 9, 7, 6, 5, 4, 3, 2, 1, 16, 17, 18, 0    ; recoded in order to make better lookup code.

; The fixed Huffman table, encoded in RLE.
; bits, count, ...  total 288 codes.
.TreeFixed  db 8, 144, 9, 112, 7, 24, 8, 8



endp


; Restores the Huffman tree from given code lengths.
; .pTree  - pointer to memory for the huffman tree structure.
; .pLengths - array of bytes, containing the lengths for every code needed.
; .count  - count of the array elements.
;
; Returns:
;   ZF=1; The tree is properly created on [.pTree]; The tree is an array of qwords
;         Every element represents one node of the binary tree. The low order dword
;         contains index to the left (0) child and the high order dword, index to
;         the right (1) child.
;         The indexes are a number of a qword element of the array xored with the index of the current element.
;         If the high order bit of the word (bit 15) is set, the word is a leaf
;         and contains NOT-ed symbol value instead of index in the array.
;
;   ZF=0  Error. The provided length list describes not valid Huffman tree.
;         In this case no values are filled in the tree.
;
; The comments in the code are according to the RFC-1951
;
proc __tree_by_lengths, .pLengths, .count, .pTree

.bl_count  rd DEFLATE_MAX_BITS + 1
.next_code rd DEFLATE_MAX_BITS + 1      ; keep immediately after .bl_count

begin
        pushad

        xor     eax, eax

        lea     edi, [.bl_count]
        mov     ecx, 2*(DEFLATE_MAX_BITS+1)
        rep stosd                               ; initializes both: .bl_count and .next_code arrays!

; 1)  Count the number of codes for each code length.  Let bl_count[N] be the number of
; codes of length N, N >= 1.

        mov     esi, [.pLengths]
        mov     ecx, [.count]

.loop:
        lodsb                               ; tree[i].len
        inc     dword [.bl_count+4*eax]     ; count one code length
        loop    .loop

; 2)  Find the numerical value of the smallest code for each code length:
; loop from 1 to [.max_bits]

;        xor     ecx, ecx         ; bits - 1    ; here ecx == 0, see the previous instr.
        xor     eax, eax         ; bl_count[0] = 0  ; but we still need to keep the original bl_count[0].
        cdq                      ; edx = 0 = code
        lea     ebx, [eax+2]     ; ebx = 2

.loop2:
        inc     ecx

        add     edx, eax                       ; code = code + bl_count[bits-1]
        shl     edx, 1                         ; code = code << 1

        mov     eax, [.bl_count+4*ecx]         ; for the next loop
        mov     [.next_code+4*ecx], edx        ; next_code[bits] = code

        sub     ebx, eax     ; data validity control.
        shl     ebx, 1

        cmp     ecx, DEFLATE_MAX_BITS
        jb      .loop2

; 3) Assign numerical values to all codes, using consecutive
;    values for all codes of the same length with the base
;    values determined at step 2. Codes that are never used
;    (which have a bit length of zero) must not be assigned a
;    value.
;
; The assigned codes are put inside the binary tree, represented in
; an array pointed by EDI

        mov     ecx, [.count]        ; count of all elements
        sub     ecx, [.bl_count]     ; minus the count of zero elements .bl_count[0]

; PKZIP exception when the Huffman tree contains single 1bit long element, the second
; is allowed to be not defined and ebx will not be zero... damn exceptions!

        cmp     ecx, 1
        jne     .check

        cmp     [.bl_count+4], 0        ; count of codes with length 1 bit.
        jne     .valid

.check:
        test    ebx, ebx
        jnz     .finish         ; it is invalid Huffman tree if EBX <> 0 here.

.valid:

; clear the tree area
        shl     ecx, 1
        mov     edi, [.pTree]   ; the needed nodes memory is (K-1)*8 bytes,
        or      eax, -1         ; where K is the count of the codes with length <> 0;
        rep stosd               ; but we allocate one element more in order to work-around the
                                ; above described exception.

        mov     edi, [.pTree]   ; the huffman tree array.
        xor     ebx, ebx        ; offset to the last free element of the tree.

        mov     esi, [.pLengths]

.loop3:
        movzx   ecx, byte [esi]     ; ecx = length of the code.
        jecxz   .next

        mov     eax, [.next_code+4*ecx]    ; next_code[len]
        inc     [.next_code+4*ecx]

        ror     eax, cl         ; start from the high order bit.
        xor     edx, edx        ; the traversing index in the tree array.

.bit_loop:
        shl     eax, 1          ; the high order bit in CF
        rcl     edx, 1          ; edx = 2*edx + CF  ; edx now points to the needed WORD element of the tree.

        dec     ecx
        jz      .put_code       ; this is the last bit, so put the code here.

.next_node:
        xor     edx, [edi+4*edx]
        jns     .bit_loop

        not     edx

        inc     ebx
        mov     [edi+4*edx], edx   ; the current free element.
        xor     [edi+4*edx], ebx   ; it is tricky here - the idea is to allow the next node offset to be obtained not by MOV, but by XOR instuction.
                                   ; for details see the tree traversing code on "bitloop1" label.
        jmp     .next_node

.put_code:
        mov     ecx, esi
        sub     ecx, [.pLengths]
        not     ecx
        xor     ecx, edx
        mov     [edi+4*edx], ecx

.next:
        inc     esi
        dec     [.count]
        jnz     .loop3

.finish:                        ; here ZF = 1 if OK and 0 if error!
        popad
        return
endp




endmodule



