1use std::collections::{BTreeMap, BTreeSet};
4
5use crate::instruction::{check_metadata, error};
6use crate::{
7 DecodedCode, InstructionError, InstructionErrorKind, InstructionId, InstructionOperand, Opcode,
8};
9
10mod scalar;
11use scalar::fixed_operand_width;
12
13pub fn encode_instructions(
19 code: &mut DecodedCode,
20 major_version: u16,
21) -> Result<Vec<u8>, InstructionError> {
22 let targets = resolve_branch_targets(code)?;
23 let offsets = layout(code, major_version)?;
24 let target_offsets: BTreeMap<_, _> = code
25 .instructions
26 .iter()
27 .zip(&offsets)
28 .map(|(located, offset)| (located.id, *offset))
29 .collect();
30 let total = code
31 .instructions
32 .iter()
33 .zip(&offsets)
34 .try_fold(0usize, |_, (located, offset)| {
35 instruction_end(&located.instruction, *offset as usize)
36 })?;
37 let mut bytes = Vec::new();
38 bytes.try_reserve_exact(total).map_err(|cause| {
39 error(
40 InstructionErrorKind::WidthOverflow,
41 total,
42 format!("cannot allocate encoded code array: {cause}"),
43 )
44 })?;
45
46 for (index, located) in code.instructions.iter().enumerate() {
47 let start = offsets[index] as usize;
48 debug_assert_eq!(bytes.len(), start);
49 encode_one(
50 &mut bytes,
51 &located.instruction,
52 start,
53 index,
54 &targets,
55 &target_offsets,
56 )?;
57 }
58
59 for (instruction_index, located) in code.instructions.iter_mut().enumerate() {
60 let start = offsets[instruction_index] as usize;
61 for (operand_index, operand) in located.instruction.operands.iter_mut().enumerate() {
62 if matches!(operand, InstructionOperand::Branch(_)) {
63 *operand = InstructionOperand::Branch(displacement(
64 start,
65 instruction_index,
66 operand_index,
67 &targets,
68 &target_offsets,
69 )?);
70 }
71 }
72 }
73
74 let mut rebuilt = BTreeMap::new();
75 for (located, offset) in code.instructions.iter_mut().zip(offsets) {
76 located.offset = offset;
77 if rebuilt.insert(offset, located.id).is_some() {
78 return Err(error(
79 InstructionErrorKind::Manifest,
80 offset as usize,
81 "duplicate encoded instruction offset",
82 ));
83 }
84 }
85 code.offsets = rebuilt;
86 Ok(bytes)
87}
88
89fn resolve_branch_targets(
90 code: &DecodedCode,
91) -> Result<BTreeMap<(usize, usize), InstructionId>, InstructionError> {
92 let ids: BTreeSet<_> = code.instructions.iter().map(|located| located.id).collect();
93 if ids.len() != code.instructions.len() {
94 return Err(error(
95 InstructionErrorKind::InvalidOperands,
96 0,
97 "instruction ids must be unique",
98 ));
99 }
100 let mut targets = BTreeMap::new();
101 for (instruction_index, located) in code.instructions.iter().enumerate() {
102 for (operand_index, operand) in located.instruction.operands.iter().enumerate() {
103 if let InstructionOperand::Branch(displacement) = operand {
104 let target = i64::from(located.offset) + i64::from(*displacement);
105 let target = u32::try_from(target).map_err(|_| {
106 error(
107 InstructionErrorKind::InvalidTarget,
108 located.offset as usize,
109 format!("branch target {target} is outside the original code array"),
110 )
111 })?;
112 let id = code.offsets.get(&target).copied().ok_or_else(|| {
113 error(
114 InstructionErrorKind::InvalidTarget,
115 located.offset as usize,
116 format!("branch target {target} is not an instruction boundary"),
117 )
118 })?;
119 if !ids.contains(&id) {
120 return Err(error(
121 InstructionErrorKind::InvalidTarget,
122 located.offset as usize,
123 format!("branch target instruction {:?} is absent", id),
124 ));
125 }
126 targets.insert((instruction_index, operand_index), id);
127 }
128 }
129 }
130 Ok(targets)
131}
132
133fn layout(code: &DecodedCode, major: u16) -> Result<Vec<u32>, InstructionError> {
134 let mut offsets = Vec::with_capacity(code.instructions.len());
135 let mut cursor = 0usize;
136 for located in &code.instructions {
137 check_metadata(located.instruction.opcode.metadata(), major, cursor)?;
138 validate_shape(&located.instruction, cursor)?;
139 offsets.push(u32::try_from(cursor).map_err(|_| {
140 error(
141 InstructionErrorKind::WidthOverflow,
142 cursor,
143 "encoded instruction offset exceeds u32",
144 )
145 })?);
146 cursor = instruction_end(&located.instruction, cursor)?;
147 }
148 u32::try_from(cursor).map_err(|_| {
149 error(
150 InstructionErrorKind::WidthOverflow,
151 cursor,
152 "encoded code array exceeds u32",
153 )
154 })?;
155 Ok(offsets)
156}
157
158fn instruction_end(
159 instruction: &crate::Instruction,
160 start: usize,
161) -> Result<usize, InstructionError> {
162 let body = if instruction.opcode == Opcode::Tableswitch {
163 let count = instruction.operands.len().checked_sub(3).ok_or_else(|| {
164 invalid(
165 start,
166 "tableswitch requires default, low, and high operands",
167 )
168 })?;
169 padding(start).checked_add(12).and_then(|size| {
170 count
171 .checked_mul(4)
172 .and_then(|entries| size.checked_add(entries))
173 })
174 } else if instruction.opcode == Opcode::Lookupswitch {
175 let tail = instruction
176 .operands
177 .len()
178 .checked_sub(1)
179 .ok_or_else(|| invalid(start, "lookupswitch requires a default operand"))?;
180 let pairs = tail / 2;
181 padding(start).checked_add(8).and_then(|size| {
182 pairs
183 .checked_mul(8)
184 .and_then(|entries| size.checked_add(entries))
185 })
186 } else {
187 fixed_operand_width(instruction, start).map(Some)?
188 }
189 .ok_or_else(|| {
190 error(
191 InstructionErrorKind::WidthOverflow,
192 start,
193 "instruction width overflows",
194 )
195 })?;
196 start
197 .checked_add(1 + usize::from(instruction.wide))
198 .and_then(|value| value.checked_add(body))
199 .ok_or_else(|| {
200 error(
201 InstructionErrorKind::WidthOverflow,
202 start,
203 "code offset overflows",
204 )
205 })
206}
207
208fn validate_shape(instruction: &crate::Instruction, start: usize) -> Result<(), InstructionError> {
209 let metadata = instruction.opcode.metadata();
210 if instruction.wide
211 && metadata.operands != "local:u1"
212 && metadata.operands != "local:u1,increment:s1"
213 {
214 return Err(error(
215 InstructionErrorKind::IllegalWide,
216 start,
217 format!("wide cannot modify {}", metadata.mnemonic),
218 ));
219 }
220 if instruction.opcode == Opcode::Tableswitch {
221 match instruction.operands.as_slice() {
222 [
223 InstructionOperand::Branch(_),
224 InstructionOperand::TableLow(low),
225 InstructionOperand::TableHigh(high),
226 branches @ ..,
227 ] if high >= low
228 && i64::from(*high) - i64::from(*low) + 1 == branches.len() as i64
229 && branches
230 .iter()
231 .all(|operand| matches!(operand, InstructionOperand::Branch(_))) =>
232 {
233 Ok(())
234 }
235 _ => Err(invalid(
236 start,
237 "tableswitch operands do not match its key range",
238 )),
239 }
240 } else if instruction.opcode == Opcode::Lookupswitch {
241 let Some((InstructionOperand::Branch(_), tail)) = instruction.operands.split_first() else {
242 return Err(invalid(start, "lookupswitch requires a default branch"));
243 };
244 if tail.len() % 2 != 0 {
245 return Err(invalid(start, "lookupswitch requires key/branch pairs"));
246 }
247 let mut previous = None;
248 for pair in tail.chunks_exact(2) {
249 let (InstructionOperand::LookupKey(key), InstructionOperand::Branch(_)) =
250 (pair[0], pair[1])
251 else {
252 return Err(invalid(start, "lookupswitch requires key/branch pairs"));
253 };
254 if previous.is_some_and(|prior| key <= prior) {
255 return Err(invalid(
256 start,
257 format!("lookupswitch key {key} is not strictly increasing"),
258 ));
259 }
260 previous = Some(key);
261 }
262 Ok(())
263 } else {
264 let expected = metadata
265 .operands
266 .split(',')
267 .filter(|field| !field.starts_with("zero:") && *field != "none")
268 .count();
269 if instruction.operands.len() != expected {
270 return Err(invalid(
271 start,
272 format!(
273 "{} expects {expected} operands, found {}",
274 metadata.mnemonic,
275 instruction.operands.len()
276 ),
277 ));
278 }
279 Ok(())
280 }
281}
282
283fn encode_one(
284 bytes: &mut Vec<u8>,
285 instruction: &crate::Instruction,
286 start: usize,
287 instruction_index: usize,
288 targets: &BTreeMap<(usize, usize), InstructionId>,
289 target_offsets: &BTreeMap<InstructionId, u32>,
290) -> Result<(), InstructionError> {
291 if instruction.wide {
292 bytes.push(Opcode::Wide as u8);
293 }
294 bytes.push(instruction.opcode as u8);
295 if instruction.opcode == Opcode::Tableswitch || instruction.opcode == Opcode::Lookupswitch {
296 bytes.resize(bytes.len() + padding(start), 0);
297 encode_switch(
298 bytes,
299 instruction,
300 start,
301 instruction_index,
302 targets,
303 target_offsets,
304 )?;
305 return Ok(());
306 }
307 let mut operand_index = 0usize;
308 for field in instruction.opcode.metadata().operands.split(',') {
309 if field == "none" {
310 continue;
311 }
312 if field == "zero:u1" {
313 bytes.push(0);
314 continue;
315 }
316 if field == "zero:u2" {
317 bytes.extend_from_slice(&[0, 0]);
318 continue;
319 }
320 let operand = instruction.operands[operand_index];
321 match (field, operand) {
322 ("value:s1" | "increment:s1", InstructionOperand::Immediate(value))
323 if !instruction.wide =>
324 {
325 push_i1(bytes, value, start)?
326 }
327 ("increment:s1", InstructionOperand::Immediate(value))
328 | ("value:s2" | "increment:s2", InstructionOperand::Immediate(value)) => {
329 push_i2(bytes, value, start)?
330 }
331 ("local:u1", InstructionOperand::Local(value)) if !instruction.wide => {
332 push_u1(bytes, value, start)?
333 }
334 ("local:u1", InstructionOperand::Local(value)) => {
335 bytes.extend_from_slice(&value.to_be_bytes())
336 }
337 ("constant_pool:u1", InstructionOperand::Constant(value)) => {
338 push_u1(bytes, value, start)?
339 }
340 ("constant_pool:u2", InstructionOperand::Constant(value)) => {
341 bytes.extend_from_slice(&value.to_be_bytes())
342 }
343 ("branch:s2", InstructionOperand::Branch(_)) => push_i2(
344 bytes,
345 displacement(
346 start,
347 instruction_index,
348 operand_index,
349 targets,
350 target_offsets,
351 )?,
352 start,
353 )?,
354 ("branch:s4", InstructionOperand::Branch(_)) => bytes.extend_from_slice(
355 &displacement(
356 start,
357 instruction_index,
358 operand_index,
359 targets,
360 target_offsets,
361 )?
362 .to_be_bytes(),
363 ),
364 ("count:u1", InstructionOperand::Count(value))
365 | ("dimensions:u1", InstructionOperand::Dimensions(value))
366 | ("atype:u1", InstructionOperand::ArrayType(value)) => bytes.push(value),
367 _ => {
368 return Err(invalid(
369 start,
370 format!("operand {operand_index} does not match manifest field {field}"),
371 ));
372 }
373 }
374 operand_index += 1;
375 }
376 Ok(())
377}
378
379fn encode_switch(
380 bytes: &mut Vec<u8>,
381 instruction: &crate::Instruction,
382 start: usize,
383 instruction_index: usize,
384 targets: &BTreeMap<(usize, usize), InstructionId>,
385 target_offsets: &BTreeMap<InstructionId, u32>,
386) -> Result<(), InstructionError> {
387 bytes.extend_from_slice(
388 &displacement(start, instruction_index, 0, targets, target_offsets)?.to_be_bytes(),
389 );
390 if instruction.opcode == Opcode::Tableswitch {
391 let InstructionOperand::TableLow(low) = instruction.operands[1] else {
392 unreachable!()
393 };
394 let InstructionOperand::TableHigh(high) = instruction.operands[2] else {
395 unreachable!()
396 };
397 bytes.extend_from_slice(&low.to_be_bytes());
398 bytes.extend_from_slice(&high.to_be_bytes());
399 for operand_index in 3..instruction.operands.len() {
400 bytes.extend_from_slice(
401 &displacement(
402 start,
403 instruction_index,
404 operand_index,
405 targets,
406 target_offsets,
407 )?
408 .to_be_bytes(),
409 );
410 }
411 } else {
412 let pairs = (instruction.operands.len() - 1) / 2;
413 bytes.extend_from_slice(&(pairs as i32).to_be_bytes());
414 for operand_index in (1..instruction.operands.len()).step_by(2) {
415 let InstructionOperand::LookupKey(key) = instruction.operands[operand_index] else {
416 unreachable!()
417 };
418 bytes.extend_from_slice(&key.to_be_bytes());
419 bytes.extend_from_slice(
420 &displacement(
421 start,
422 instruction_index,
423 operand_index + 1,
424 targets,
425 target_offsets,
426 )?
427 .to_be_bytes(),
428 );
429 }
430 }
431 Ok(())
432}
433
434fn displacement(
435 start: usize,
436 instruction_index: usize,
437 operand_index: usize,
438 targets: &BTreeMap<(usize, usize), InstructionId>,
439 target_offsets: &BTreeMap<InstructionId, u32>,
440) -> Result<i32, InstructionError> {
441 let target = targets
442 .get(&(instruction_index, operand_index))
443 .expect("validated branch target");
444 let target_offset = target_offsets.get(target).ok_or_else(|| {
445 error(
446 InstructionErrorKind::InvalidTarget,
447 start,
448 format!("branch target {:?} is absent", target),
449 )
450 })?;
451 i32::try_from(i64::from(*target_offset) - start as i64).map_err(|_| {
452 error(
453 InstructionErrorKind::WidthOverflow,
454 start,
455 "branch displacement exceeds s4",
456 )
457 })
458}
459
460fn padding(start: usize) -> usize {
461 (4 - ((start + 1) % 4)) % 4
462}
463fn push_u1(bytes: &mut Vec<u8>, value: u16, start: usize) -> Result<(), InstructionError> {
464 bytes.push(u8::try_from(value).map_err(|_| {
465 error(
466 InstructionErrorKind::WidthOverflow,
467 start,
468 format!("unsigned operand {value} exceeds u1"),
469 )
470 })?);
471 Ok(())
472}
473fn push_i1(bytes: &mut Vec<u8>, value: i32, start: usize) -> Result<(), InstructionError> {
474 bytes.push(i8::try_from(value).map_err(|_| {
475 error(
476 InstructionErrorKind::WidthOverflow,
477 start,
478 format!("signed operand {value} exceeds s1"),
479 )
480 })? as u8);
481 Ok(())
482}
483fn push_i2(bytes: &mut Vec<u8>, value: i32, start: usize) -> Result<(), InstructionError> {
484 bytes.extend_from_slice(
485 &i16::try_from(value)
486 .map_err(|_| {
487 error(
488 InstructionErrorKind::WidthOverflow,
489 start,
490 format!("signed operand {value} exceeds s2"),
491 )
492 })?
493 .to_be_bytes(),
494 );
495 Ok(())
496}
497fn invalid(offset: usize, message: impl Into<String>) -> InstructionError {
498 error(InstructionErrorKind::InvalidOperands, offset, message)
499}