1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92

// Copyright 2018 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::io;
use std::io::Write;

use crate::builder::MerkleTreeBuilder;
use crate::tree::MerkleTree;

/// A `MerkleTreeWriter` wraps a [`MerkleTreeBuilder`] and another type that implements
/// [`std::io::Write`].
///
/// `MerkleTreeWriter`s can be used to compute a [`MerkleTree`] while streaming data from one
/// location to another. To simply compute a [`MerkleTree`] without chaining writes to a separate
/// Writer, see [`MerkleTreeBuilder`].
///
/// # Examples
/// ```
/// # use fuchsia_merkle::*;
/// # use std::io::{Result,Write};
/// # fn main() -> Result<()> {
/// let data = vec![0xff; 8192];
/// let mut written = Vec::new();
/// {
///     let mut writer = MerkleTreeWriter::new(&mut written);
///     writer.write_all(&data)?;
///     assert_eq!(
///         writer.finish().root(),
///         "68d131bc271f9c192d4f6dcd8fe61bef90004856da19d0f2f514a7f4098b0737"
///             .parse()
///             .unwrap()
///     );
/// }
/// assert_eq!(written, data);
/// #     Ok(())
/// # }
/// ```
pub struct MerkleTreeWriter<W> {
    inner: W,
    builder: MerkleTreeBuilder,
}

impl<W: Write> MerkleTreeWriter<W> {
    /// Creates a new `MerkleTreeWriter`
    pub fn new(inner: W) -> Self {
        MerkleTreeWriter { inner, builder: MerkleTreeBuilder::new() }
    }

    /// Finalizes all levels of the merkle tree, converting this `MerkleTreeWriter` instance into a
    /// [`MerkleTree`].
    pub fn finish(self) -> MerkleTree {
        self.builder.finish()
    }
}

impl<W: Write> Write for MerkleTreeWriter<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.inner.write_all(buf)?;
        self.builder.write(buf);
        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        self.inner.flush()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::Hash;
    use test_case::test_case;

    #[allow(clippy::unused_unit)]
    #[test_case(vec![], "15ec7bf0b50732b49f8228e07d24365338f9e3ab994b00af08e5a3bffe55fd8b" ; "test_empty")]
    #[test_case(vec![0xFF; 8192], "68d131bc271f9c192d4f6dcd8fe61bef90004856da19d0f2f514a7f4098b0737"; "test_oneblock")]
    #[test_case(vec![0xFF; 65536], "f75f59a944d2433bc6830ec243bfefa457704d2aed12f30539cd4f18bf1d62cf"; "test_small")]
    #[test_case(vec![0xFF; 2105344], "7d75dfb18bfd48e03b5be4e8e9aeea2f89880cb81c1551df855e0d0a0cc59a67"; "test_large")]
    #[test_case(vec![0xFF; 2109440], "7577266aa98ce587922fdc668c186e27f3c742fb1b732737153b70ae46973e43"; "test_unaligned")]
    fn tests(input: Vec<u8>, output: &str) {
        let mut written = Vec::with_capacity(input.len());
        let actual = {
            let mut tree = MerkleTreeWriter::new(&mut written);
            tree.write_all(input.as_slice()).unwrap();
            tree.finish().root()
        };
        let expected: Hash = output.parse().unwrap();
        assert_eq!(expected, actual);
        assert_eq!(input, written);
    }
}