From e21196f8b0f6b5a8f2922dd62a180a92bc798be8 Mon Sep 17 00:00:00 2001
From: chyyran <ronny@ronnychan.ca>
Date: Tue, 9 Aug 2022 03:44:21 -0400
Subject: [PATCH] Alternative API with only a single generic argument for
 BitTree

---
 src/decode/lzma.rs       |  4 ++--
 src/decode/rangecoder.rs | 31 +++++++++----------------------
 src/encode/rangecoder.rs |  4 ++--
 3 files changed, 13 insertions(+), 26 deletions(-)

diff --git a/src/decode/lzma.rs b/src/decode/lzma.rs
index f468841..dbcc89a 100644
--- a/src/decode/lzma.rs
+++ b/src/decode/lzma.rs
@@ -167,8 +167,8 @@ pub(crate) struct DecoderState {
     pub(crate) lzma_props: LzmaProperties,
     unpacked_size: Option<u64>,
     literal_probs: Vec2D<u16>,
-    pos_slot_decoder: [BitTree<6, { bittree_probs_len::<6>() }>; 4],
-    align_decoder: BitTree<4, { bittree_probs_len::<4>() }>,
+    pos_slot_decoder: [BitTree<{ bittree_probs_len::<6>() }>; 4],
+    align_decoder: BitTree<{ bittree_probs_len::<4>() }>,
     pos_decoders: [u16; 115],
     is_match: [u16; 192], // true = LZ, false = literal
     is_rep: [u16; 12],
diff --git a/src/decode/rangecoder.rs b/src/decode/rangecoder.rs
index b7b356f..c1fa7a2 100644
--- a/src/decode/rangecoder.rs
+++ b/src/decode/rangecoder.rs
@@ -150,44 +150,31 @@ where
     }
 }
 
-/// macro for compile-time const assertions
-macro_rules! const_assert {
-    ($($list:ident : $ty:ty),* => $expr:expr) => {{
-        struct Assert<$(const $list: $ty,)*>;
-        impl<$(const $list: $ty,)*> Assert<$($list,)*> {
-            const OK: u8 = 0 - !($expr) as u8;
-        }
-        Assert::<$($list,)*>::OK
-    }};
-    ($expr:expr) => {
-        const OK: u8 = 0 - !($expr) as u8;
-    };
-}
-
 // const fn helper to parameterize the length of the bittree probability array.
 pub const fn bittree_probs_len<const NUM_BITS: usize>() -> usize {
     1 << NUM_BITS
 }
 
 #[derive(Debug, Clone)]
-pub struct BitTree<const NUM_BITS: usize, const PROBS_ARRAY_LEN: usize> {
+pub struct BitTree<const PROBS_ARRAY_LEN: usize> {
     probs: [u16; PROBS_ARRAY_LEN],
 }
 
-impl<const NUM_BITS: usize, const PROBS_ARRAY_LEN: usize> BitTree<NUM_BITS, PROBS_ARRAY_LEN> {
+impl<const PROBS_ARRAY_LEN: usize> BitTree<PROBS_ARRAY_LEN> {
     pub fn new() -> Self {
-        const_assert!(NUM_BITS: usize, PROBS_ARRAY_LEN: usize => PROBS_ARRAY_LEN == bittree_probs_len::<NUM_BITS>());
         BitTree {
             probs: [0x400; PROBS_ARRAY_LEN],
         }
     }
 
+    const NUM_BITS: usize = PROBS_ARRAY_LEN.trailing_zeros() as usize;
+
     pub fn parse<R: io::BufRead>(
         &mut self,
         rangecoder: &mut RangeDecoder<R>,
         update: bool,
     ) -> io::Result<u32> {
-        rangecoder.parse_bit_tree(NUM_BITS, &mut self.probs, update)
+        rangecoder.parse_bit_tree(Self::NUM_BITS, &mut self.probs, update)
     }
 
     pub fn parse_reverse<R: io::BufRead>(
@@ -195,7 +182,7 @@ impl<const NUM_BITS: usize, const PROBS_ARRAY_LEN: usize> BitTree<NUM_BITS, PROB
         rangecoder: &mut RangeDecoder<R>,
         update: bool,
     ) -> io::Result<u32> {
-        rangecoder.parse_reverse_bit_tree(NUM_BITS, &mut self.probs, 0, update)
+        rangecoder.parse_reverse_bit_tree(Self::NUM_BITS, &mut self.probs, 0, update)
     }
 }
 
@@ -203,9 +190,9 @@ impl<const NUM_BITS: usize, const PROBS_ARRAY_LEN: usize> BitTree<NUM_BITS, PROB
 pub struct LenDecoder {
     choice: u16,
     choice2: u16,
-    low_coder: [BitTree<3, { bittree_probs_len::<3>() }>; 16],
-    mid_coder: [BitTree<3, { bittree_probs_len::<3>() }>; 16],
-    high_coder: BitTree<8, { bittree_probs_len::<8>() }>,
+    low_coder: [BitTree<{ bittree_probs_len::<3>() }>; 16],
+    mid_coder: [BitTree<{ bittree_probs_len::<3>() }>; 16],
+    high_coder: BitTree<{ bittree_probs_len::<8>() }>,
 }
 
 impl LenDecoder {
diff --git a/src/encode/rangecoder.rs b/src/encode/rangecoder.rs
index e02deca..49b5a14 100644
--- a/src/encode/rangecoder.rs
+++ b/src/encode/rangecoder.rs
@@ -266,7 +266,7 @@ mod test {
 
         let mut bufread = BufReader::new(buf.as_slice());
         let mut decoder = RangeDecoder::new(&mut bufread).unwrap();
-        let mut tree = decode::rangecoder::BitTree::<NUM_BITS, PROBS_LEN>::new();
+        let mut tree = decode::rangecoder::BitTree::<PROBS_LEN>::new();
         for &v in values {
             assert_eq!(tree.parse(&mut decoder, true).unwrap(), v);
         }
@@ -313,7 +313,7 @@ mod test {
 
         let mut bufread = BufReader::new(buf.as_slice());
         let mut decoder = RangeDecoder::new(&mut bufread).unwrap();
-        let mut tree = decode::rangecoder::BitTree::<NUM_BITS, PROBS_LEN>::new();
+        let mut tree = decode::rangecoder::BitTree::<PROBS_LEN>::new();
         for &v in values {
             assert_eq!(tree.parse_reverse(&mut decoder, true).unwrap(), v);
         }