openvm_native_compiler/asm/
code.rs

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
use alloc::{collections::BTreeMap, format};
use core::{fmt, fmt::Display};

use openvm_circuit::arch::instructions::instruction::DebugInfo;
use openvm_stark_backend::p3_field::{ExtensionField, PrimeField32};

use super::AsmInstruction;

/// A basic block of assembly instructions.
#[derive(Debug, Clone, Default)]
pub struct BasicBlock<F, EF>(
    pub(crate) Vec<AsmInstruction<F, EF>>,
    pub(crate) Vec<Option<DebugInfo>>,
);

impl<F: PrimeField32, EF: ExtensionField<F>> BasicBlock<F, EF> {
    /// Creates a new basic block.
    pub fn new() -> Self {
        Self(Vec::new(), Vec::new())
    }

    /// Pushes an instruction to a basic block.
    pub(crate) fn push(
        &mut self,
        instruction: AsmInstruction<F, EF>,
        debug_info: Option<DebugInfo>,
    ) {
        self.0.push(instruction);
        self.1.push(debug_info);
    }
}

/// Assembly code for a program.
pub struct AssemblyCode<F, EF> {
    pub blocks: Vec<BasicBlock<F, EF>>,
    pub labels: BTreeMap<F, String>,
}

impl<F: PrimeField32, EF: ExtensionField<F>> AssemblyCode<F, EF> {
    /// Creates a new assembly code.
    pub fn new(blocks: Vec<BasicBlock<F, EF>>, labels: BTreeMap<F, String>) -> Self {
        Self { blocks, labels }
    }

    pub fn size(&self) -> usize {
        self.blocks.iter().map(|block| block.0.len()).sum()
    }
}

impl<F: PrimeField32, EF: ExtensionField<F>> Display for AssemblyCode<F, EF> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (i, block) in self.blocks.iter().enumerate() {
            writeln!(
                f,
                "{}:",
                self.labels
                    .get(&F::from_canonical_u32(i as u32))
                    .unwrap_or(&format!(".L{}", i))
            )?;
            for instruction in &block.0 {
                write!(f, "        ")?;
                instruction.fmt(&self.labels, f)?;
                writeln!(f)?;
            }
        }
        Ok(())
    }
}