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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
use crate::{Ecc, Endian, Error, Result, BE, LE, NE};
use byteorder::{ByteOrder, ReadBytesExt, WriteBytesExt};
use std::io::{Read, Write};

/// A simple hierarchical storage system for strings.
#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
pub struct Hierarchical {
    /// The primary string key for this level of the hierarchy.
    key: String,
    /// The content of this level.
    content: Vec<String>,
    /// The optional children.
    children: Vec<Hierarchical>,
}

impl Hierarchical {
    /// Ecc identifier type.
    const ID: Ecc = Ecc::new("STR_HIER");

    /// Create a new hierarchical structure.
    pub fn new<T: Into<String>>(key: T, content: Vec<String>, children: Vec<Self>) -> Self {
        Self {
            key: key.into(),
            content,
            children,
        }
    }

    /// Get the key value.
    pub fn key(&self) -> &str {
        &self.key
    }

    /// Get the key mutably.
    pub fn key_mut(&mut self) -> &mut String {
        &mut self.key
    }

    /// Get the content.
    pub fn content(&self) -> &[String] {
        &self.content
    }

    /// Get the content mutably.
    pub fn content_mut(&mut self) -> &mut Vec<String> {
        &mut self.content
    }

    /// Push a new content item.
    pub fn push<T: Into<String>>(&mut self, item: T) {
        self.content.push(item.into());
    }

    /// Get the children.
    pub fn children(&self) -> &[Self] {
        &self.children
    }

    /// Get the children mutably.
    pub fn children_mut(&mut self) -> &mut Vec<Self> {
        &mut self.children
    }

    /// Push a child hierarchy.
    pub fn push_child(&mut self, item: Self) {
        self.children.push(item);
    }

    /// Convert the hierarhcy into a byte buffer.
    pub fn to_bytes<E: ByteOrder>(self) -> Result<Vec<u8>> {
        // Create the output buffer to write to.
        let mut bytes = vec![];
        let writer: &mut dyn Write = &mut bytes;

        // Write the identification ID.
        Self::ID.write::<E>(writer)?;

        // Write the structure.
        self.write_structure::<E>(writer)?;

        // Bytes is now filled and nothing went wrong.
        // Likely the only real issues possible are running out of memory
        // due to incorrect content or strings > 65k.
        Ok(bytes)
    }

    /// Write a hierarchical structure recursively.
    fn write_structure<E: ByteOrder>(self, writer: &mut dyn Write) -> Result<()> {
        // Write the key value.
        Self::write_string::<E>(&self.key, writer)?;

        // Write the count of content strings.
        writer.write_u32::<E>(self.content.len() as u32)?;
        // Write each content string.
        for s in self.content {
            Self::write_string::<E>(&s, writer)?;
        }

        // Write each child.
        // Write the count of children.
        writer.write_u32::<E>(self.children.len() as u32)?;
        for c in self.children {
            c.write_structure::<E>(writer)?;
        }

        Ok(())
    }

    // Write a string into the given stream.
    fn write_string<E: ByteOrder>(value: &str, writer: &mut dyn Write) -> Result<()> {
        // Write the key value.
        if value.len() > core::u16::MAX as usize {
            return Err(Error::Invalid("String length greater than max u16!".into()));
        }
        writer.write_u16::<E>(value.len() as u16)?;
        writer.write_all(value.as_bytes())?;

        Ok(())
    }

    /// Create a hierarchical structure from the given bytes.
    /// TODO: Should probably just give this a reader rather than taking
    /// the bytes directly..  Hmmm
    pub fn from_bytes(mut bytes: &[u8]) -> Result<Self> {
        // Create a reader.
        let reader: &mut dyn Read = &mut bytes;

        // Detect endian by reading in local endian.
        let id = Ecc::from(reader.read_u64::<NE>()?);
        match id.endian(Self::ID) {
            Some(Endian::Big) => Ok(Self::from_reader::<BE>(reader)?),
            Some(Endian::Little) => Ok(Self::from_reader::<LE>(reader)?),
            None => Err(Error::Invalid("Not a valid hierarchical.".into())),
        }
    }

    // Read the hierarchy with the given endian.
    fn from_reader<E: ByteOrder>(reader: &mut dyn Read) -> Result<Self> {
        // Read the key.
        let key = Self::read_string::<E>(reader)?;

        // Read the content.
        let content_count = reader.read_u32::<E>()?;
        let mut content = vec![];
        for _ in 0..content_count {
            content.push(Self::read_string::<E>(reader)?);
        }

        // Read the children.
        let child_count = reader.read_u32::<E>()?;
        let mut children = vec![];
        for _ in 0..child_count {
            children.push(Self::from_reader::<E>(reader)?);
        }

        Ok(Self {
            key,
            content,
            children,
        })
    }

    // Read a string from the given reader.
    fn read_string<E: ByteOrder>(reader: &mut dyn Read) -> Result<String> {
        let len = reader.read_u16::<E>()?;
        let mut buffer = vec![0; len as usize];
        reader.read_exact(buffer.as_mut_slice())?;

        Ok(std::str::from_utf8(buffer.as_slice())?.into())
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_serialization() {
        let test = Hierarchical::new(
            "Bah",
            vec![String::from("1"), String::from("2")],
            vec![Hierarchical::new(
                "Humbug",
                vec![String::from("3"), String::from("4")],
                vec![],
            )],
        );

        let bytes = test.clone().to_bytes::<NE>().unwrap();
        let result = Hierarchical::from_bytes(bytes.as_slice()).unwrap();
        assert_eq!(test, result);
    }
}