alloc: Reduce unnecessary Vec allocations

Also changes literal_probs array from a 2D Vec<Vec<u16>> to a Grid backed by a contiguous Vec<u16>
SnowflakePowered · Aug 6, 2022 · ecc71a0 · ecc71a0
1 parent bf45a8a
commit ecc71a0
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 11 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -16,6 +16,7 @@ byteorder = "^1.0.0"
 crc = "^1.0.0"
 log = { version = "^0.4.14", optional = true }
 env_logger = { version = "^0.8.3", optional = true }
+grid = { version = "0.6.1-alpha.0",  git = "https://github.com/chyyran/grid-rs", branch = "feature-fill"}
 
 [dev-dependencies]
 rust-lzma = "0.5"

diff --git a/src/decode/lzma.rs b/src/decode/lzma.rs
@@ -5,6 +5,7 @@ use crate::decompress::Options;
 use crate::decompress::UnpackedSize;
 use crate::error;
 use byteorder::{LittleEndian, ReadBytesExt};
+use grid::Grid;
 use std::io;
 
 /// Maximum input data that can be processed in one iteration.
@@ -167,8 +168,8 @@ pub(crate) struct DecoderState {
     partial_input_buf: std::io::Cursor<[u8; MAX_REQUIRED_INPUT]>,
     pub(crate) lzma_props: LzmaProperties,
     unpacked_size: Option<u64>,
-    literal_probs: Vec<Vec<u16>>,
-    pos_slot_decoder: Vec<rangecoder::BitTree>,
+    literal_probs: Grid<u16>,
+    pos_slot_decoder: [rangecoder::BitTree; 4],
     align_decoder: rangecoder::BitTree,
     pos_decoders: [u16; 115],
     is_match: [u16; 192], // true = LZ, false = literal
@@ -190,8 +191,13 @@ impl DecoderState {
             partial_input_buf: std::io::Cursor::new([0; MAX_REQUIRED_INPUT]),
             lzma_props,
             unpacked_size,
-            literal_probs: vec![vec![0x400; 0x300]; 1 << (lzma_props.lc + lzma_props.lp)],
-            pos_slot_decoder: vec![rangecoder::BitTree::new(6); 4],
+            literal_probs: Grid::init(1 << (lzma_props.lc + lzma_props.lp), 0x300, 0x400),
+            pos_slot_decoder: [
+                rangecoder::BitTree::new(6),
+                rangecoder::BitTree::new(6),
+                rangecoder::BitTree::new(6),
+                rangecoder::BitTree::new(6),
+            ],
             align_decoder: rangecoder::BitTree::new(4),
             pos_decoders: [0x400; 115],
             is_match: [0x400; 192],
@@ -210,11 +216,10 @@ impl DecoderState {
     pub fn reset_state(&mut self, new_props: LzmaProperties) {
         new_props.validate();
         if self.lzma_props.lc + self.lzma_props.lp == new_props.lc + new_props.lp {
-            // We can reset here by filling the existing buffer with 0x400.
-            self.literal_probs.iter_mut().for_each(|v| v.fill(0x400))
+            self.literal_probs.fill(0x400);
         } else {
             // We need to reallocate because of the new size of `lc+lp`.
-            self.literal_probs = vec![vec![0x400; 0x300]; 1 << (new_props.lc + new_props.lp)];
+            self.literal_probs = Grid::init(1 << (new_props.lc + new_props.lp), 0x300, 0x400);
         }
 
         self.lzma_props = new_props;

diff --git a/src/decode/rangecoder.rs b/src/decode/rangecoder.rs
@@ -190,8 +190,8 @@ impl BitTree {
 pub struct LenDecoder {
     choice: u16,
     choice2: u16,
-    low_coder: Vec<BitTree>,
-    mid_coder: Vec<BitTree>,
+    low_coder: [BitTree; 16],
+    mid_coder: [BitTree; 16],
     high_coder: BitTree,
 }
 
@@ -200,8 +200,14 @@ impl LenDecoder {
         LenDecoder {
             choice: 0x400,
             choice2: 0x400,
-            low_coder: vec![BitTree::new(3); 16],
-            mid_coder: vec![BitTree::new(3); 16],
+            low_coder: [BitTree::new(3), BitTree::new(3), BitTree::new(3), BitTree::new(3),
+                        BitTree::new(3), BitTree::new(3), BitTree::new(3), BitTree::new(3),
+                        BitTree::new(3), BitTree::new(3), BitTree::new(3), BitTree::new(3),
+                        BitTree::new(3), BitTree::new(3), BitTree::new(3), BitTree::new(3)],
+            mid_coder: [BitTree::new(3), BitTree::new(3), BitTree::new(3), BitTree::new(3),
+                        BitTree::new(3), BitTree::new(3), BitTree::new(3), BitTree::new(3),
+                        BitTree::new(3), BitTree::new(3), BitTree::new(3), BitTree::new(3),
+                        BitTree::new(3), BitTree::new(3), BitTree::new(3), BitTree::new(3)],
             high_coder: BitTree::new(8),
         }
     }