1use crate::asm::{opcode::ParseOp, ToBytes, ToOpcode, TryFromBytes};
4
5#[derive(Clone, Debug, PartialEq)]
20pub struct BytecodeMapped<Op, Bytes = Vec<u8>> {
21    bytecode: Bytes,
23    op_indices: Vec<usize>,
27    _op_ty: core::marker::PhantomData<Op>,
29}
30
31#[derive(Clone, Copy, Debug, PartialEq)]
33pub struct BytecodeMappedSlice<'a, Op> {
34    bytecode: &'a [u8],
36    op_indices: &'a [usize],
38    _op_ty: core::marker::PhantomData<Op>,
40}
41
42impl<Op> BytecodeMapped<Op, Vec<u8>> {
43    pub fn push_op(&mut self, op: Op)
45    where
46        Op: ToBytes,
47    {
48        self.op_indices.push(self.bytecode.len());
49        self.bytecode.extend(op.to_bytes());
50    }
51}
52
53impl<Op, Bytes> BytecodeMapped<Op, Bytes>
54where
55    Bytes: core::ops::Deref<Target = [u8]>,
56{
57    pub fn try_from_bytes(bytes: Bytes) -> Result<BytecodeMapped<Op, Bytes>, Op::Error>
62    where
63        Op: ToOpcode + TryFromBytes,
64        Op::Opcode: ParseOp<Op = Op> + TryFrom<u8>,
65        Op::Error: From<<Op::Opcode as TryFrom<u8>>::Error> + From<<Op::Opcode as ParseOp>::Error>,
66    {
67        let bytecode = bytes.deref();
68        let mut op_indices = Vec::with_capacity(bytecode.len() / std::mem::size_of::<Op>());
69        let mut iter_enum = bytecode.iter().enumerate();
70        while let Some((ix, &opcode_byte)) = iter_enum.next() {
71            let opcode = Op::Opcode::try_from(opcode_byte)?;
72            let mut op_bytes = iter_enum.by_ref().map(|(_, &byte)| byte);
73            let _op = opcode.parse_op(&mut op_bytes)?;
74            op_indices.push(ix);
75        }
76        Ok(BytecodeMapped {
77            bytecode: bytes,
78            op_indices,
79            _op_ty: core::marker::PhantomData,
80        })
81    }
82
83    pub fn as_slice(&self) -> BytecodeMappedSlice<Op> {
85        BytecodeMappedSlice {
86            bytecode: self.bytecode(),
87            op_indices: self.op_indices(),
88            _op_ty: self._op_ty,
89        }
90    }
91
92    pub fn bytecode(&self) -> &[u8] {
94        self.bytecode.deref()
95    }
96
97    pub fn ops_from(&self, start: usize) -> Option<BytecodeMappedSlice<Op>> {
103        Some(BytecodeMappedSlice {
104            bytecode: self.bytecode(),
105            op_indices: self.op_indices.get(start..)?,
106            _op_ty: self._op_ty,
107        })
108    }
109
110    pub fn op(&self, ix: usize) -> Option<Op>
112    where
113        Op: TryFromBytes,
114    {
115        let slice = self.ops_from(ix)?;
116        slice.ops().next()
117    }
118
119    pub fn ops(&self) -> impl '_ + Iterator<Item = Op>
121    where
122        Op: TryFromBytes,
123    {
124        expect_ops_from_indices(self.bytecode(), self.op_indices.iter().copied())
125    }
126}
127
128impl<Op, Bytes> BytecodeMapped<Op, Bytes> {
129    pub fn op_indices(&self) -> &[usize] {
131        &self.op_indices
132    }
133}
134
135impl<'a, Op> BytecodeMappedSlice<'a, Op> {
136    pub fn op_indices(self) -> &'a [usize] {
138        self.op_indices
139    }
140
141    pub fn ops(self) -> impl 'a + Iterator<Item = Op>
143    where
144        Op: TryFromBytes,
145    {
146        expect_ops_from_indices(self.bytecode, self.op_indices.iter().copied())
147    }
148}
149
150impl<Op> Default for BytecodeMapped<Op> {
153    fn default() -> Self {
154        BytecodeMapped {
155            bytecode: Default::default(),
156            op_indices: Default::default(),
157            _op_ty: Default::default(),
158        }
159    }
160}
161
162impl<Op> FromIterator<Op> for BytecodeMapped<Op>
164where
165    Op: ToBytes,
166{
167    fn from_iter<T: IntoIterator<Item = Op>>(iter: T) -> Self {
168        let iter = iter.into_iter();
169        let (min, _) = iter.size_hint();
170        let mut mapped = BytecodeMapped {
171            bytecode: Vec::with_capacity(min),
172            op_indices: Vec::with_capacity(min),
173            _op_ty: core::marker::PhantomData,
174        };
175        iter.for_each(|op| mapped.push_op(op));
176        mapped
177    }
178}
179
180impl<Op> TryFrom<Vec<u8>> for BytecodeMapped<Op>
182where
183    Op: ToOpcode + TryFromBytes,
184    Op::Opcode: ParseOp<Op = Op> + TryFrom<u8>,
185    Op::Error: From<<Op::Opcode as TryFrom<u8>>::Error> + From<<Op::Opcode as ParseOp>::Error>,
186{
187    type Error = Op::Error;
188    fn try_from(bytecode: Vec<u8>) -> Result<Self, Self::Error> {
189        Self::try_from_bytes(bytecode)
190    }
191}
192
193impl<'a, Op> TryFrom<&'a [u8]> for BytecodeMapped<Op, &'a [u8]>
195where
196    Op: ToOpcode + TryFromBytes,
197    Op::Opcode: ParseOp<Op = Op> + TryFrom<u8>,
198    Op::Error: From<<Op::Opcode as TryFrom<u8>>::Error> + From<<Op::Opcode as ParseOp>::Error>,
199{
200    type Error = Op::Error;
201    fn try_from(bytecode: &'a [u8]) -> Result<Self, Self::Error> {
202        Self::try_from_bytes(bytecode)
203    }
204}
205
206fn expect_ops_from_indices<'a, Op>(
209    bytecode: &'a [u8],
210    op_indices: impl 'a + IntoIterator<Item = usize>,
211) -> impl 'a + Iterator<Item = Op>
212where
213    Op: TryFromBytes,
214{
215    const EXPECT_MSG: &str = "validated upon construction";
216    op_indices.into_iter().map(|ix| {
217        let mut bytes = bytecode[ix..].iter().copied();
218        Op::try_from_bytes(&mut bytes)
219            .expect(EXPECT_MSG)
220            .expect(EXPECT_MSG)
221    })
222}