From bf0e1f71c0bdc15018c6c01647e8ad8f5851d891 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kevin=20L=C3=A4ufer?= <laeufer@cs.berkeley.edu>
Date: Thu, 9 Nov 2023 11:15:36 -0500
Subject: [PATCH] wavemem: correctly re-code 4-state to 2-state

---
 src/fst.rs          |  9 ++++++---
 src/vcd.rs          |  5 +++--
 src/wavemem.rs      | 46 +++++++++++++++++++++++++++++++++++++++++----
 tests/diff_tests.rs | 12 +++++++++++-
 4 files changed, 62 insertions(+), 10 deletions(-)

diff --git a/src/fst.rs b/src/fst.rs
index 8c98a65..6c29419 100644
--- a/src/fst.rs
+++ b/src/fst.rs
@@ -50,13 +50,16 @@ impl<R: BufRead + Seek> SignalSource for FstWaveDatabase<R> {
         let mut index_and_time = time_table.next().unwrap();
 
         // store signals
-        let mut signals = Vec::with_capacity(ids.len());
-        let foo = |time: u64, _handle: FstSignalHandle, _value: FstSignalValue| {
+        let mut signals = Vec::new();
+        let idx_to_pos: HashMap<usize, usize> =
+            HashMap::from_iter(ids.iter().map(|(r, _)| r.index()).enumerate());
+        let foo = |time: u64, handle: FstSignalHandle, _value: FstSignalValue| {
             // determine time index
             while *(index_and_time.1) < time {
                 index_and_time = time_table.next().unwrap();
             }
-            let _time_idx = index_and_time.0;
+            let time_idx = index_and_time.0;
+            let signal_pos = idx_to_pos[&handle.get_index()];
             todo!()
         };
 
diff --git a/src/vcd.rs b/src/vcd.rs
index 0d752ef..4b7b048 100644
--- a/src/vcd.rs
+++ b/src/vcd.rs
@@ -263,6 +263,7 @@ where
         .join(", ")
 }
 
+#[derive(Debug, PartialEq)]
 enum VcdCmd {
     Date,
     Timescale,
@@ -793,14 +794,14 @@ x%i"
         let mut buf = Vec::with_capacity(128);
         let input_0 = b"$upscope $end";
         let (cmd_0, body_0) = read_command(&mut input_0.as_slice(), &mut buf).unwrap();
-        assert_eq!(cmd_0, b"upscope");
+        assert_eq!(cmd_0, VcdCmd::UpScope);
         assert!(body_0.is_empty());
 
         // test with more whitespace
         buf.clear();
         let input_1 = b" \t $upscope \n $end  \n ";
         let (cmd_1, body_1) = read_command(&mut input_1.as_slice(), &mut buf).unwrap();
-        assert_eq!(cmd_1, b"upscope");
+        assert_eq!(cmd_1, VcdCmd::UpScope);
         assert!(body_1.is_empty());
     }
 }
diff --git a/src/wavemem.rs b/src/wavemem.rs
index 8e74995..099f05a 100644
--- a/src/wavemem.rs
+++ b/src/wavemem.rs
@@ -224,12 +224,11 @@ fn load_fixed_len_signal(
             other => {
                 let num_bytes = usize_div_ceil(other as usize, 4);
                 let mut buf = vec![0u8; num_bytes];
+                data.read_exact(&mut buf.as_mut()).unwrap();
                 if is_two_state {
-                    todo!()
-                } else {
-                    data.read_exact(&mut buf.as_mut()).unwrap();
-                    out.append(&mut buf);
+                    four_state_to_two_state(&mut buf);
                 }
+                out.append(&mut buf);
                 //
                 time_idx_delta_raw
             }
@@ -241,6 +240,37 @@ fn load_fixed_len_signal(
     (out, time_indices)
 }
 
+#[inline]
+fn four_state_to_two_state(buf: &mut Vec<u8>) {
+    let result_len = usize_div_ceil(buf.len(), 2);
+    let is_uneven = result_len * 2 > buf.len();
+    for ii in 0..result_len {
+        let (msb, lsb) = if is_uneven {
+            if ii == 0 {
+                (0u8, buf[0])
+            } else {
+                (buf[ii * 2 - 1], buf[ii * 2])
+            }
+        } else {
+            (buf[ii * 2], buf[ii * 2 + 1])
+        };
+        buf[ii] = compress_four_state_to_two_state(msb, lsb);
+    }
+    buf.truncate(result_len);
+}
+
+#[inline]
+fn compress_four_state_to_two_state(msb: u8, lsb: u8) -> u8 {
+    (msb & (1 << 6)) << 1 | // msb[6] -> out[7]
+    (msb & (1 << 4)) << 2 | // msb[4] -> out[6]
+    (msb & (1 << 2)) << 3 | // msb[2] -> out[5]
+    (msb & (1 << 0)) << 4 | // msb[0] -> out[4]
+    (lsb & (1 << 6)) >> 3 | // lsb[6] -> out[3]
+    (lsb & (1 << 4)) >> 2 | // lsb[4] -> out[2]
+    (lsb & (1 << 2)) >> 1 | // lsb[2] -> out[1]
+    (lsb & (1 << 0)) >> 0 // lsb[0] -> out[0]
+}
+
 #[inline]
 fn load_signal_strings(data: &mut impl Read, time_idx_offset: u32) -> (Vec<String>, Vec<u32>) {
     let mut out = Vec::new();
@@ -796,4 +826,12 @@ mod tests {
             } // great!
         }
     }
+
+    #[test]
+    fn test_four_state_to_two_state() {
+        let mut input0 = vec![0b01010001u8, 0b00010100u8, 0b00010000u8];
+        let expected0 = [0b1101u8, 0b01100100u8];
+        four_state_to_two_state(&mut input0);
+        assert_eq!(input0, expected0);
+    }
 }
diff --git a/tests/diff_tests.rs b/tests/diff_tests.rs
index d46fdc9..dd59336 100644
--- a/tests/diff_tests.rs
+++ b/tests/diff_tests.rs
@@ -121,7 +121,17 @@ fn diff_hierarchy_item(ref_item: &vcd::ScopeItem, our_item: HierarchyItem, our_h
     }
 }
 
-fn diff_signals<R: BufRead>(_ref_reader: &mut vcd::Parser<R>, _our: &mut Waveform) {
+fn diff_signals<R: BufRead>(ref_reader: &mut vcd::Parser<R>, our: &mut Waveform) {
+    // load all signals
+    let all_signals: Vec<_> = our
+        .hierarchy()
+        .get_unique_signals_vars()
+        .iter()
+        .flatten()
+        .map(|v| v.signal_idx())
+        .collect();
+    our.load_signals(&all_signals);
+
     println!("TODO")
 }