compiler: reorg intermediate representation (#66)

* compiler: reorg intermediate representation

* cleanup

Author: tcoratger

crates/lean_compiler/src/a_simplify_lang.rs

@@ -1,6 +1,6 @@
 use crate::{
     Counter, F,
-    intermediate_bytecode::HighLevelOperation,
+    ir::HighLevelOperation,
     lang::{
         Boolean, ConstExpression, ConstMallocLabel, Expression, Function, Line, Program,
         SimpleExpr, Var,

crates/lean_compiler/src/b_compile_intermediate.rs

@@ -1,4 +1,4 @@
-use crate::{F, a_simplify_lang::*, intermediate_bytecode::*, lang::*, precompiles::*};
+use crate::{F, a_simplify_lang::*, ir::*, lang::*, precompiles::*};
 use lean_vm::*;
 use p3_field::Field;
 use std::{

crates/lean_compiler/src/c_compile_final.rs

@@ -1,4 +1,4 @@
-use crate::{F, PUBLIC_INPUT_START, ZERO_VEC_PTR, intermediate_bytecode::*, lang::*};
+use crate::{F, PUBLIC_INPUT_START, ZERO_VEC_PTR, ir::*, lang::*};
 use lean_vm::*;
 use p3_field::{PrimeCharacteristicRing, PrimeField32};
 use std::collections::BTreeMap;

crates/lean_compiler/src/ir/bytecode.rs

@@ -0,0 +1,51 @@
+use super::instruction::IntermediateInstruction;
+use lean_vm::Label;
+use std::collections::BTreeMap;
+use std::fmt::{Display, Formatter};
+
+/// Container for the complete intermediate representation of a program.
+///
+/// This structure holds all the compiled intermediate bytecode along with
+/// metadata needed for execution and analysis.
+#[derive(Debug, Clone)]
+pub struct IntermediateBytecode {
+    /// Main bytecode organized by function labels.
+    ///
+    /// Each label corresponds to a function entry point.
+    pub bytecode: BTreeMap<Label, Vec<IntermediateInstruction>>,
+
+    /// Match statement bytecode blocks.
+    ///
+    /// Each match statement produces multiple case blocks.
+    pub match_blocks: Vec<Vec<Vec<IntermediateInstruction>>>,
+
+    /// Memory requirements for each function.
+    ///
+    /// Maps function names to their stack frame size.
+    pub memory_size_per_function: BTreeMap<String, usize>,
+}
+
+impl Display for IntermediateBytecode {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        for (label, instructions) in &self.bytecode {
+            writeln!(f, "\n{label}:")?;
+            for instruction in instructions {
+                writeln!(f, "  {instruction}")?;
+            }
+        }
+        for (i, match_blocks) in self.match_blocks.iter().enumerate() {
+            writeln!(f, "\nMatch {i}:")?;
+            for (j, case) in match_blocks.iter().enumerate() {
+                writeln!(f, "  Case {j}:")?;
+                for instruction in case {
+                    writeln!(f, "    {instruction}")?;
+                }
+            }
+        }
+        writeln!(f, "\nMemory size per function:")?;
+        for (function_name, size) in &self.memory_size_per_function {
+            writeln!(f, "{function_name}: {size}")?;
+        }
+        Ok(())
+    }
+}

…an_compiler/src/intermediate_bytecode.rs …ates/lean_compiler/src/ir/instruction.rscrates/lean_compiler/src/intermediate_bytecode.rs renamed to crates/lean_compiler/src/ir/instruction.rs crates/lean_compiler/src/intermediate_bytecode.rs renamed to crates/lean_compiler/src/ir/instruction.rs

@@ -1,85 +1,10 @@
-use std::collections::BTreeMap;
-use std::fmt::Display;
-use std::fmt::Formatter;
-
-use p3_field::PrimeCharacteristicRing;
-use p3_field::PrimeField64;
-use utils::ToUsize;
-
-use crate::{F, lang::ConstExpression};
-use lean_vm::*;
-
-#[derive(Debug, Clone)]
-pub struct IntermediateBytecode {
-    pub bytecode: BTreeMap<Label, Vec<IntermediateInstruction>>,
-    pub match_blocks: Vec<Vec<Vec<IntermediateInstruction>>>,
-    pub memory_size_per_function: BTreeMap<String, usize>,
-}
-
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub enum IntermediateValue {
-    Constant(ConstExpression),
-    Fp,
-    MemoryAfterFp { offset: ConstExpression }, // m[fp + offset]
-}
-
-impl From<ConstExpression> for IntermediateValue {
-    fn from(value: ConstExpression) -> Self {
-        Self::Constant(value)
-    }
-}
-impl TryFrom<HighLevelOperation> for Operation {
-    type Error = String;
-
-    fn try_from(value: HighLevelOperation) -> Result<Self, Self::Error> {
-        match value {
-            HighLevelOperation::Add => Ok(Self::Add),
-            HighLevelOperation::Mul => Ok(Self::Mul),
-            _ => Err(format!("Cannot convert {value:?} to +/x")),
-        }
-    }
-}
-
-#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
-pub enum IntermediaryMemOrFpOrConstant {
-    MemoryAfterFp { offset: ConstExpression }, // m[fp + offset]
-    Fp,
-    Constant(ConstExpression),
-}
-
-impl IntermediateValue {
-    pub const fn label(label: Label) -> Self {
-        Self::Constant(ConstExpression::label(label))
-    }
-
-    pub const fn is_constant(&self) -> bool {
-        matches!(self, Self::Constant(_))
-    }
-}
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
-pub enum HighLevelOperation {
-    Add,
-    Mul,
-    Sub,
-    Div, // in the end everything compiles to either Add or Mul
-    Exp, // Exponentiation, only for const expressions
-    Mod, // Modulo, only for const expressions
-}
-
-impl HighLevelOperation {
-    pub fn eval(&self, a: F, b: F) -> F {
-        match self {
-            Self::Add => a + b,
-            Self::Mul => a * b,
-            Self::Sub => a - b,
-            Self::Div => a / b,
-            Self::Exp => a.exp_u64(b.as_canonical_u64()),
-            Self::Mod => F::from_usize(a.to_usize() % b.to_usize()),
-        }
-    }
-}
+use super::operation::HighLevelOperation;
+use super::value::{IntermediaryMemOrFpOrConstant, IntermediateValue};
+use crate::lang::ConstExpression;
+use lean_vm::{Operation, SourceLineNumber};
+use std::fmt::{Display, Formatter};
 
+/// Core instruction type for the intermediate representation.
 #[derive(Debug, Clone)]
 pub enum IntermediateInstruction {
     Computation {
@@ -200,64 +125,9 @@ impl IntermediateInstruction {
     }
 }
 
-impl Display for IntermediateValue {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        match self {
-            Self::Constant(value) => write!(f, "{value}"),
-            Self::Fp => write!(f, "fp"),
-            Self::MemoryAfterFp { offset } => {
-                write!(f, "m[fp + {offset}]")
-            }
-        }
-    }
-}
-
-impl Display for IntermediaryMemOrFpOrConstant {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        match self {
-            Self::MemoryAfterFp { offset } => write!(f, "m[fp + {offset}]"),
-            Self::Fp => write!(f, "fp"),
-            Self::Constant(c) => write!(f, "{c}"),
-        }
-    }
-}
-
 impl Display for IntermediateInstruction {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
-            Self::Deref {
-                shift_0,
-                shift_1,
-                res,
-            } => write!(f, "{res} = m[m[fp + {shift_0}] + {shift_1}]"),
-            Self::DotProduct {
-                arg0,
-                arg1,
-                res,
-                size,
-            } => write!(f, "dot_product({arg0}, {arg1}, {res}, {size})"),
-            Self::MultilinearEval {
-                coeffs,
-                point,
-                res,
-                n_vars,
-            } => write!(f, "multilinear_eval({coeffs}, {point}, {res}, {n_vars})"),
-            Self::DecomposeBits {
-                res_offset,
-                to_decompose,
-            } => {
-                write!(f, "m[fp + {res_offset}..] = decompose_bits(")?;
-                for (i, expr) in to_decompose.iter().enumerate() {
-                    if i > 0 {
-                        write!(f, ", ")?;
-                    }
-                    write!(f, "{expr}")?;
-                }
-                write!(f, ")")
-            }
-            Self::CounterHint { res_offset } => {
-                write!(f, "m[fp + {res_offset}] = counter_hint()")
-            }
             Self::Computation {
                 operation,
                 arg_a,
@@ -266,6 +136,11 @@ impl Display for IntermediateInstruction {
             } => {
                 write!(f, "{res} = {arg_a} {operation} {arg_c}")
             }
+            Self::Deref {
+                shift_0,
+                shift_1,
+                res,
+            } => write!(f, "{res} = m[m[fp + {shift_0}] + {shift_1}]"),
             Self::Panic => write!(f, "panic"),
             Self::Jump { dest, updated_fp } => {
                 if let Some(fp) = updated_fp {
@@ -291,6 +166,18 @@ impl Display for IntermediateInstruction {
             Self::Poseidon2_24 { arg_a, arg_b, res } => {
                 write!(f, "{res} = poseidon2_24({arg_a}, {arg_b})")
             }
+            Self::DotProduct {
+                arg0,
+                arg1,
+                res,
+                size,
+            } => write!(f, "dot_product({arg0}, {arg1}, {res}, {size})"),
+            Self::MultilinearEval {
+                coeffs,
+                point,
+                res,
+                n_vars,
+            } => write!(f, "multilinear_eval({coeffs}, {point}, {res}, {n_vars})"),
             Self::Inverse { arg, res_offset } => {
                 write!(f, "m[fp + {res_offset}] = inverse({arg})")
             }
@@ -309,6 +196,22 @@ impl Display for IntermediateInstruction {
                     write!(f, "m[fp + {offset}] = request_memory({size})")
                 }
             }
+            Self::DecomposeBits {
+                res_offset,
+                to_decompose,
+            } => {
+                write!(f, "m[fp + {res_offset}..] = decompose_bits(")?;
+                for (i, expr) in to_decompose.iter().enumerate() {
+                    if i > 0 {
+                        write!(f, ", ")?;
+                    }
+                    write!(f, "{expr}")?;
+                }
+                write!(f, ")")
+            }
+            Self::CounterHint { res_offset } => {
+                write!(f, "m[fp + {res_offset}] = counter_hint()")
+            }
             Self::Print { line_info, content } => {
                 write!(f, "print {line_info}: ")?;
                 for (i, c) in content.iter().enumerate() {
@@ -323,41 +226,3 @@ impl Display for IntermediateInstruction {
         }
     }
 }
-
-impl Display for HighLevelOperation {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        match self {
-            Self::Add => write!(f, "+"),
-            Self::Mul => write!(f, "*"),
-            Self::Sub => write!(f, "-"),
-            Self::Div => write!(f, "/"),
-            Self::Exp => write!(f, "**"),
-            Self::Mod => write!(f, "%"),
-        }
-    }
-}
-
-impl Display for IntermediateBytecode {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        for (label, instructions) in &self.bytecode {
-            writeln!(f, "\n{label}:")?;
-            for instruction in instructions {
-                writeln!(f, "  {instruction}")?;
-            }
-        }
-        for (i, match_blocks) in self.match_blocks.iter().enumerate() {
-            writeln!(f, "\nMatch {i}:")?;
-            for (j, case) in match_blocks.iter().enumerate() {
-                writeln!(f, "  Case {j}:")?;
-                for instruction in case {
-                    writeln!(f, "    {instruction}")?;
-                }
-            }
-        }
-        writeln!(f, "\nMemory size per function:")?;
-        for (function_name, size) in &self.memory_size_per_function {
-            writeln!(f, "{function_name}: {size}")?;
-        }
-        Ok(())
-    }
-}

crates/lean_compiler/src/ir/mod.rs

@@ -0,0 +1,11 @@
+//! Intermediate Representation (IR) for the Lean compiler.
+
+pub mod bytecode;
+pub mod instruction;
+pub mod operation;
+pub mod value;
+
+pub use bytecode::IntermediateBytecode;
+pub use instruction::IntermediateInstruction;
+pub use operation::HighLevelOperation;
+pub use value::{IntermediaryMemOrFpOrConstant, IntermediateValue};