11package io .arsha .api .lib ;
22
3- import io .arsha .api .exceptions .CannotBeRegisteredException ;
43import jakarta .annotation .Nullable ;
54import jakarta .validation .constraints .NotNull ;
65import java .io .ByteArrayInputStream ;
1413import org .thshsh .struct .Struct ;
1514
1615// Implementation from https://github.com/shrddr/huffman_heap
16+ // This isn't entirely accurate anymore, but it used to basically be a direct copy so im keeping it
1717@ Component
1818public class HuffmanDecoder {
1919
20- record Node (@ Nullable String character , @ NotNull Long frequency , @ Nullable Node left , @ Nullable Node right ) {
21- public static Node fromEntry (Map .Entry <String , Long > entry ) {
22- return new Node (entry .getKey (), entry .getValue (), null , null );
23- }
24-
25- public static Node branch (Node left , Node right ) {
26- return new Node (null , left .frequency () + right .frequency (), left , right );
27- }
28- }
29-
3020 private static Map <String , Long > getFrequencies (InputStream buffer ) throws IOException {
3121 var data = Struct .create ("<3I" ).unpack (buffer );
3222 var chars = (long ) data .get (2 );
@@ -35,7 +25,7 @@ private static Map<String, Long> getFrequencies(InputStream buffer) throws IOExc
3525 for (int i = 0 ; i < chars ; i ++) {
3626 var charData = Struct .create ("<Ic3b" ).unpack (buffer );
3727 var frequency = (long ) charData .get (0 );
38- var character = new String (new byte [] { (byte ) charData .get (1 ) });
28+ var character = new String (new byte []{ (byte ) charData .get (1 )});
3929 frequencyMap .put (character , frequency );
4030 }
4131 return frequencyMap ;
@@ -45,9 +35,7 @@ private static Node getRoot(InputStream stream) throws IOException {
4535 var frequencies = getFrequencies (stream );
4636
4737 var heap = new PriorityQueue <>(Comparator .comparingLong (Node ::frequency ));
48- frequencies .entrySet ().stream ()
49- .map (Node ::fromEntry )
50- .forEach (heap ::add );
38+ frequencies .entrySet ().stream ().map (Node ::fromEntry ).forEach (heap ::add );
5139
5240 while (heap .size () > 1 ) {
5341 var left = heap .poll ();
@@ -62,40 +50,50 @@ private static Node getRoot(InputStream stream) throws IOException {
6250 return heap .peek ();
6351 }
6452
65- private static String unpack (Node root , StringBuilder packed ) {
53+ public String unpack (byte [] buffer ) throws IOException {
54+ var stream = new ByteArrayInputStream (buffer );
55+
56+ var root = getRoot (stream );
57+ var sizes = Struct .create ("<3I" ).unpack (stream );
58+ var bits = (long ) sizes .getFirst (); // The bit size of the encoded data
59+ // var bytes = (long) sizes.get(1); // The byte size of the encoded data
60+ var expectedBytes = (long ) sizes .get (2 ); // The expected byte size of the unpacked data
61+
62+ var counter = 0 ;
6663 var current = root ;
6764 var unpacked = new StringBuilder ();
68- for (var i = 0 ; i < packed .length (); i ++) {
69- current = packed .charAt (i ) == '0' ? current .left () : current .right ();
70- assert current != null ;
65+ outer :
66+ for (var b : stream .readAllBytes ()) {
67+ if (unpacked .length () > expectedBytes ) {
68+ throw new IllegalStateException ("Unpacked data exceeds expected size" );
69+ }
7170
72- if (current .left () == null && current .right () == null ) {
73- unpacked .append (current .character ());
74- current = root ;
71+ for (int i = 7 ; i >= 0 ; i --) {
72+ if (counter >= bits ) {
73+ break outer ;
74+ }
75+ int bit = (b >> i ) & 1 ;
76+ current = bit == 0 ? current .left () : current .right ();
77+ assert current != null ;
78+ if (current .left () == null && current .right () == null ) {
79+ unpacked .append (current .character ());
80+ current = root ;
81+ }
82+ counter ++;
7583 }
7684 }
85+
7786 return unpacked .toString ();
7887 }
7988
80- public String decode (byte [] buffer ) throws CannotBeRegisteredException , IOException {
81- var testResult = new String (buffer );
82- if (testResult .contains ("resultMsg" )) { // Probably an item that can't be registered on market (code 8)
83- throw new CannotBeRegisteredException (testResult );
84- }
85-
86- var dataIn = new ByteArrayInputStream (buffer );
87- var root = getRoot (dataIn );
88-
89- var data = Struct .create ("<3I" ).unpack (dataIn );
90- var bits = (long ) data .getFirst (); // The bit size of the encoded data
91- // var bytes = (long) data.get(1); // The byte size of the encoded data
92- // var expectedSize = (long) data.get(2); // The expected byte size of the unpacked data
89+ record Node (@ Nullable String character , @ NotNull Long frequency , @ Nullable Node left , @ Nullable Node right ) {
9390
94- var bin = new StringBuilder ();
95- for (byte b : dataIn .readAllBytes ())
96- bin .append (Integer .toBinaryString (b & 255 | 256 ).substring (1 ));
97- if (bin .length () > bits ) bin .setLength (Math .toIntExact (bits ));
91+ public static Node fromEntry (Map .Entry <String , Long > entry ) {
92+ return new Node (entry .getKey (), entry .getValue (), null , null );
93+ }
9894
99- return unpack (root , bin );
95+ public static Node branch (Node left , Node right ) {
96+ return new Node (null , left .frequency () + right .frequency (), left , right );
97+ }
10098 }
10199}
0 commit comments