1
use crate::entity::{EntityRef, Split, Tag, Transaction};
2
use crate::error::{Error, Result};
3
use crate::host;
4
use scripting_format::{
5
    ContextType, ENTITY_HEADER_SIZE, EntityHeader, EntityType, GLOBAL_HEADER_SIZE, GlobalHeader,
6
};
7

            
8
pub struct InputReader {
9
    buffer: &'static [u8],
10
    header: GlobalHeader,
11
    base_offset: u32,
12
}
13

            
14
impl InputReader {
15
    pub fn new() -> Result<Self> {
16
        let base_offset = host::input_offset();
17
        let base = base_offset as *const u8;
18

            
19
        // SAFETY: WASM linear memory is valid for the lifetime of the module.
20
        // The host guarantees the input buffer is properly initialized.
21
        let (buffer, header) = unsafe {
22
            let header_slice = core::slice::from_raw_parts(base, GLOBAL_HEADER_SIZE);
23
            let header = GlobalHeader::from_bytes(header_slice).ok_or(Error::InvalidMagic)?;
24
            let total_size =
25
                (header.strings_pool_offset - base_offset + header.strings_pool_size) as usize;
26
            let buffer = core::slice::from_raw_parts(base, total_size);
27
            (buffer, header)
28
        };
29

            
30
        Ok(Self {
31
            buffer,
32
            header,
33
            base_offset,
34
        })
35
    }
36

            
37
    pub fn context_type(&self) -> Result<ContextType> {
38
        ContextType::try_from(self.header.context_type).map_err(|()| Error::InvalidEntityType)
39
    }
40

            
41
    pub fn primary_entity_type(&self) -> Result<EntityType> {
42
        EntityType::try_from(self.header.primary_entity_type).map_err(|()| Error::InvalidEntityType)
43
    }
44

            
45
    #[must_use]
46
    pub fn entity_count(&self) -> u32 {
47
        self.header.input_entity_count
48
    }
49

            
50
    #[must_use]
51
    pub fn primary_entity_idx(&self) -> u32 {
52
        self.header.primary_entity_idx
53
    }
54

            
55
    fn strings_pool(&self) -> &[u8] {
56
        let offset = (self.header.strings_pool_offset - self.base_offset) as usize;
57
        let size = self.header.strings_pool_size as usize;
58
        &self.buffer[offset..offset + size]
59
    }
60

            
61
    fn entity_offset(&self, idx: u32) -> Result<usize> {
62
        if idx >= self.header.input_entity_count {
63
            return Err(Error::EntityNotFound);
64
        }
65
        let mut offset = (self.header.entities_offset - self.base_offset) as usize;
66
        for _ in 0..idx {
67
            let header_bytes = &self.buffer[offset..offset + ENTITY_HEADER_SIZE];
68
            let header = EntityHeader::from_bytes(header_bytes).ok_or(Error::InvalidHeader)?;
69
            offset += ENTITY_HEADER_SIZE + header.data_size as usize;
70
        }
71
        Ok(offset)
72
    }
73

            
74
    pub fn entity(&self, idx: u32) -> Result<EntityRef<'_>> {
75
        let offset = self.entity_offset(idx)?;
76
        let header_bytes = &self.buffer[offset..offset + ENTITY_HEADER_SIZE];
77
        let header = EntityHeader::from_bytes(header_bytes).ok_or(Error::InvalidHeader)?;
78
        let data_offset = offset + ENTITY_HEADER_SIZE;
79
        let data_size = header.data_size as usize;
80
        let data = &self.buffer[data_offset..data_offset + data_size];
81

            
82
        Ok(EntityRef {
83
            header,
84
            data,
85
            strings_pool: self.strings_pool(),
86
        })
87
    }
88

            
89
    pub fn primary(&self) -> Result<EntityRef<'_>> {
90
        self.entity(self.header.primary_entity_idx)
91
    }
92

            
93
    pub fn primary_transaction(&self) -> Result<Transaction> {
94
        self.primary()?.as_transaction()
95
    }
96

            
97
    pub fn transaction(&self, idx: u32) -> Result<Transaction> {
98
        self.entity(idx)?.as_transaction()
99
    }
100

            
101
    pub fn split(&self, idx: u32) -> Result<Split> {
102
        self.entity(idx)?.as_split()
103
    }
104

            
105
    pub fn tag(&self, idx: u32) -> Result<Tag<'_>> {
106
        self.entity(idx)?.as_tag()
107
    }
108

            
109
    pub fn splits_for(&self, parent_idx: u32) -> impl Iterator<Item = Result<Split>> + '_ {
110
        let count = self.header.input_entity_count;
111
        (0..count).filter_map(move |idx| {
112
            let entity = match self.entity(idx) {
113
                Ok(e) => e,
114
                Err(e) => return Some(Err(e)),
115
            };
116
            if entity.header.parent_idx == parent_idx as i32 {
117
                match entity.entity_type() {
118
                    Ok(EntityType::Split) => Some(entity.as_split()),
119
                    Ok(_) => None,
120
                    Err(e) => Some(Err(e)),
121
                }
122
            } else {
123
                None
124
            }
125
        })
126
    }
127

            
128
    pub fn tags_for(&self, parent_idx: u32) -> impl Iterator<Item = Result<Tag<'_>>> + '_ {
129
        let count = self.header.input_entity_count;
130
        (0..count).filter_map(move |idx| {
131
            let entity = match self.entity(idx) {
132
                Ok(e) => e,
133
                Err(e) => return Some(Err(e)),
134
            };
135
            if entity.header.parent_idx == parent_idx as i32 {
136
                match entity.entity_type() {
137
                    Ok(EntityType::Tag) => Some(entity.as_tag()),
138
                    Ok(_) => None,
139
                    Err(e) => Some(Err(e)),
140
                }
141
            } else {
142
                None
143
            }
144
        })
145
    }
146
}