openvm_native_compiler/asm/
code.rs

1use alloc::{collections::BTreeMap, format};
2use core::{fmt, fmt::Display};
3
4use openvm_circuit::arch::instructions::instruction::DebugInfo;
5use openvm_stark_backend::p3_field::{ExtensionField, PrimeField32};
6
7use super::AsmInstruction;
8
9/// A basic block of assembly instructions.
10#[derive(Debug, Clone, Default)]
11pub struct BasicBlock<F, EF>(
12    pub(crate) Vec<AsmInstruction<F, EF>>,
13    pub(crate) Vec<Option<DebugInfo>>,
14);
15
16impl<F: PrimeField32, EF: ExtensionField<F>> BasicBlock<F, EF> {
17    /// Creates a new basic block.
18    pub fn new() -> Self {
19        Self(Vec::new(), Vec::new())
20    }
21
22    /// Pushes an instruction to a basic block.
23    pub(crate) fn push(
24        &mut self,
25        instruction: AsmInstruction<F, EF>,
26        debug_info: Option<DebugInfo>,
27    ) {
28        self.0.push(instruction);
29        self.1.push(debug_info);
30    }
31}
32
33/// Assembly code for a program.
34pub struct AssemblyCode<F, EF> {
35    pub blocks: Vec<BasicBlock<F, EF>>,
36    pub labels: BTreeMap<F, String>,
37}
38
39impl<F: PrimeField32, EF: ExtensionField<F>> AssemblyCode<F, EF> {
40    /// Creates a new assembly code.
41    pub fn new(blocks: Vec<BasicBlock<F, EF>>, labels: BTreeMap<F, String>) -> Self {
42        Self { blocks, labels }
43    }
44
45    pub fn size(&self) -> usize {
46        self.blocks.iter().map(|block| block.0.len()).sum()
47    }
48}
49
50impl<F: PrimeField32, EF: ExtensionField<F>> Display for AssemblyCode<F, EF> {
51    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
52        for (i, block) in self.blocks.iter().enumerate() {
53            writeln!(
54                f,
55                "{}:",
56                self.labels
57                    .get(&F::from_canonical_u32(i as u32))
58                    .unwrap_or(&format!(".L{}", i))
59            )?;
60            for instruction in &block.0 {
61                write!(f, "        ")?;
62                instruction.fmt(&self.labels, f)?;
63                writeln!(f)?;
64            }
65        }
66        Ok(())
67    }
68}