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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

use bao_tree::{
    io::{fsm::BaoContentItem, sync::WriteAt},
    BaoTree,
};
use bytes::Bytes;

use crate::{
    util::{compute_outboard, copy_limited_slice, SparseMemFile},
    IROH_BLOCK_SIZE,
};

/// Mutable in memory storage for a bao file.
///
/// This is used for incomplete files if they are not big enough to warrant
/// writing to disk. We must keep track of ranges in both data and outboard
/// that have been written to, and track the most precise known size.
#[derive(Debug, Default)]
pub struct MutableMemStorage {
    /// Data file, can be any size.
    pub data: SparseMemFile,
    /// Outboard file, must be a multiple of 64 bytes.
    pub outboard: SparseMemFile,
    /// Size that was announced as we wrote that chunk
    pub sizes: SizeInfo,
}

/// Keep track of the most precise size we know of.
///
/// When in memory, we don't have to write the size for every chunk to a separate
/// slot, but can just keep the best one.
#[derive(Debug, Default)]
pub struct SizeInfo {
    pub offset: u64,
    pub size: u64,
}

impl SizeInfo {
    /// Create a new size info for a complete file of size `size`.
    pub(crate) fn complete(size: u64) -> Self {
        let mask = (1 << IROH_BLOCK_SIZE.chunk_log()) - 1;
        // offset of the last bao chunk in a file of size `size`
        let last_chunk_offset = size & mask;
        Self {
            offset: last_chunk_offset,
            size,
        }
    }

    /// Write a size at the given offset. The size at the highest offset is going to be kept.
    fn write(&mut self, offset: u64, size: u64) {
        // >= instead of > because we want to be able to update size 0, the initial value.
        if offset >= self.offset {
            self.offset = offset;
            self.size = size;
        }
    }

    /// The current size, representing the most correct size we know.
    pub fn current_size(&self) -> u64 {
        self.size
    }
}

impl MutableMemStorage {
    /// Create a new mutable mem storage from the given data
    pub fn complete(
        bytes: Bytes,
        cb: impl Fn(u64) + Send + Sync + 'static,
    ) -> (Self, iroh_base::hash::Hash) {
        let (hash, outboard) = compute_outboard(&bytes[..], bytes.len() as u64, move |offset| {
            cb(offset);
            Ok(())
        })
        .unwrap();
        let outboard = outboard.unwrap_or_default();
        let res = Self {
            data: bytes.to_vec().into(),
            outboard: outboard.into(),
            sizes: SizeInfo::complete(bytes.len() as u64),
        };
        (res, hash)
    }

    pub(super) fn current_size(&self) -> u64 {
        self.sizes.current_size()
    }

    pub(super) fn read_data_at(&self, offset: u64, len: usize) -> Bytes {
        copy_limited_slice(&self.data, offset, len)
    }

    pub(super) fn data_len(&self) -> u64 {
        self.data.len() as u64
    }

    pub(super) fn read_outboard_at(&self, offset: u64, len: usize) -> Bytes {
        copy_limited_slice(&self.outboard, offset, len)
    }

    pub(super) fn outboard_len(&self) -> u64 {
        self.outboard.len() as u64
    }

    pub(super) fn write_batch(
        &mut self,
        size: u64,
        batch: &[BaoContentItem],
    ) -> std::io::Result<()> {
        let tree = BaoTree::new(size, IROH_BLOCK_SIZE);
        for item in batch {
            match item {
                BaoContentItem::Parent(parent) => {
                    if let Some(offset) = tree.pre_order_offset(parent.node) {
                        let o0 = offset
                            .checked_mul(64)
                            .expect("u64 overflow multiplying to hash pair offset");
                        let o1 = o0.checked_add(32).expect("u64 overflow");
                        let outboard = &mut self.outboard;
                        outboard.write_all_at(o0, parent.pair.0.as_bytes().as_slice())?;
                        outboard.write_all_at(o1, parent.pair.1.as_bytes().as_slice())?;
                    }
                }
                BaoContentItem::Leaf(leaf) => {
                    self.sizes.write(leaf.offset, size);
                    self.data.write_all_at(leaf.offset, leaf.data.as_ref())?;
                }
            }
        }
        Ok(())
    }
}