From 05e6edcfc54deed588366bf3b3c1bc9289b296a1 Mon Sep 17 00:00:00 2001
From: ChillerDragon <ChillerDragon@gmail.com>
Date: Fri, 2 Aug 2024 14:40:31 +0800
Subject: [PATCH] huffman set_bits_r not sure if it works

---
 src/huffman/construct_tree.asm |  15 +++++
 src/huffman/set_bits_r.asm     | 111 +++++++++++++++++++++++++++++++++
 src/huffman_decompress.asm     |   0
 src/macros.asm                 |  40 ++++++++++++
 src/teeworlds_asmr.asm         |   3 +-
 tests/assert.asm               |   3 +-
 6 files changed, 170 insertions(+), 2 deletions(-)
 create mode 100644 src/huffman/construct_tree.asm
 create mode 100644 src/huffman/set_bits_r.asm
 delete mode 100644 src/huffman_decompress.asm

diff --git a/src/huffman/construct_tree.asm b/src/huffman/construct_tree.asm
new file mode 100644
index 0000000..a5e6dd7
--- /dev/null
+++ b/src/huffman/construct_tree.asm
@@ -0,0 +1,15 @@
+; ; m_aNodes array of CNODE structs
+; huff_nodes resb HUFF_CNODE_SIZE * HUFFMAN_MAX_NODES
+; 
+; ; m_apDecodeLut array of CNODE structs
+; huff_decode_lut resb HUFFMAN_LUTSIZE * 8
+; 
+; ; pointer to start node
+; huff_start_node resb 8
+; 
+; ; integer with amount of nodes
+; huff_num_nodes resb 4
+
+_huff_construct_tree:
+    ret
+
diff --git a/src/huffman/set_bits_r.asm b/src/huffman/set_bits_r.asm
new file mode 100644
index 0000000..dfaeed0
--- /dev/null
+++ b/src/huffman/set_bits_r.asm
@@ -0,0 +1,111 @@
+; ; m_aNodes array of CNODE structs
+; huff_nodes resb HUFF_CNODE_SIZE * HUFFMAN_MAX_NODES
+; 
+; ; m_apDecodeLut array of CNODE structs
+; huff_decode_lut resb HUFFMAN_LUTSIZE * 8
+; 
+; ; pointer to start node
+; huff_start_node resb 8
+; 
+; ; integer with amount of nodes
+; huff_num_nodes resb 4
+
+_huff_setbits_r:
+    ; _huff_setbits_r [rax] [rdi] [rsi]
+    ;  rax = *pNode
+    ;  rdi = int Bits
+    ;  rsi = unsigned Depth
+    push_registers
+
+    ; leaf1
+    push rdi
+    push rsi
+
+    mov rcx, 0
+    mov word cx, [rax + HUFF_CNODE_LEAF_1_OFFSET]
+    cmp cx, 0xffff
+    je ._huff_setbits_r_no_leaf1_recursion
+
+    ; &m_aNodes[pNode->m_aLeafs[1]]
+    lea rax, [huff_nodes + rcx]
+
+    ; bits | (1<<Depth)
+    mov r9, 1
+    shift_left r9, rsi
+    xor rdi, r9
+    shl edx,cl
+
+    ; depth + 1
+    inc rsi
+    call _huff_setbits_r
+
+    pop rsi
+    pop rdi
+
+._huff_setbits_r_no_leaf1_recursion:
+
+    ; leaf0
+    mov word cx, [rax + HUFF_CNODE_LEAF_0_OFFSET]
+    cmp cx, 0xffff
+    je ._huff_setbits_r_no_leaf0_recursion
+
+    ; &m_aNodes[pNode->m_aLeafs[0]]
+    lea rax, [huff_nodes + rcx]
+
+    ; bits are still in rdi
+
+    ; depth + 1
+    inc rsi
+    call _huff_setbits_r
+
+._huff_setbits_r_no_leaf0_recursion:
+
+    ; num bits
+    mov dword ecx, [rax + HUFF_CNODE_LEAF_0_OFFSET]
+    je ._huff_setbits_r_got_num_bits
+    jmp ._huff_setbits_r_end
+
+._huff_setbits_r_got_num_bits:
+    ; node.bits = bits
+    lea rcx, [rax + HUFF_CNODE_BITS_OFFSET]
+    mov dword [rcx], edi
+
+    ; node.num_bits = depth
+    lea rcx, [rax + HUFF_CNODE_NUM_BITS_OFFSET]
+    mov dword [rcx], esi
+
+._huff_setbits_r_end:
+
+    pop_registers
+    ret
+
+%macro huff_setbits_r 3
+    ; huff_setbits_r [*pNode] [int Bits] [unsigned Depth]
+    push rax
+    push rdi
+    push rsi
+
+    mov rbp, rsp
+
+    ; *pNode = 8 byte
+    ; int Bits = 4 byte
+    ; unsigned Depth = 4 byte
+    ; 8 + 4 + 4 = 16
+    sub rsp, 16
+
+    mov qword [rbp-16], %1
+    mov dword [rbp-8], %2
+    mov dword [rbp-4], %3
+
+    mov rax, [rbp-16]
+    mov rdi, [rbp-8]
+    mov rsi, [rbp-4]
+    call _huff_setbits_r
+
+    mov rsp, rbp
+
+    pop rsi
+    pop rdi
+    pop rax
+%endmacro
+
diff --git a/src/huffman_decompress.asm b/src/huffman_decompress.asm
deleted file mode 100644
index e69de29..0000000
diff --git a/src/macros.asm b/src/macros.asm
index 9aa92ad..bba7158 100644
--- a/src/macros.asm
+++ b/src/macros.asm
@@ -190,3 +190,43 @@
     pop rcx
     pop rbx
 %endmacro
+
+%macro shift_left 2
+    ; shift_left [register] [amount]
+    ;
+    ; I saw with objdump that `shl edx, cl` is possible.
+    ; But nasm refuses to compile it
+    ;
+    ; So this macro polyfills it
+    ; shift left n bits
+    mov rbp, rsp
+    sub rsp, 8
+
+    push rcx
+    push r12
+
+    ; register
+    mov [rbp-8], %1
+
+    ; amount
+    mov r12, %2
+
+    mov rcx, 0
+    %%shift_left_loop:
+        inc rcx
+        cmp rcx, r12
+        jge %%shift_left_loop_end
+
+        shl qword [rbp-8], 1
+
+        jmp %%shift_left_loop
+    %%shift_left_loop_end:
+
+    pop r12
+    pop rcx
+
+    mov %1, [rbp-8]
+
+    mov rsp, rbp
+%endmacro
+
diff --git a/src/teeworlds_asmr.asm b/src/teeworlds_asmr.asm
index 2b032b4..77ec7b4 100644
--- a/src/teeworlds_asmr.asm
+++ b/src/teeworlds_asmr.asm
@@ -139,7 +139,8 @@ section .text
 %include "src/on_game.asm"
 %include "src/packet_packer.asm"
 %include "src/pack_int.asm"
-%include "src/huffman_decompress.asm"
+%include "src/huffman/set_bits_r.asm"
+%include "src/huffman/construct_tree.asm"
 
 print_udp:
     print s_got_udp
diff --git a/tests/assert.asm b/tests/assert.asm
index d1fee1b..5e218c7 100644
--- a/tests/assert.asm
+++ b/tests/assert.asm
@@ -61,7 +61,8 @@ section .text
 %include "src/on_game.asm"
 %include "src/packet_packer.asm"
 %include "src/pack_int.asm"
-%include "src/huffman_decompress.asm"
+%include "src/huffman/set_bits_r.asm"
+%include "src/huffman/construct_tree.asm"
 
 assert_ok:
     push_registers