MassaLabs: Implement a Intermediate Representation to improve the compilation process

Cargo.toml

@@ -4,6 +4,7 @@ members = [
     "parser",
     "pass",
     "ir",
+    "codegen/air",
     "codegen/masm",
     "codegen/winterfell",
 ]

air-script/Cargo.toml

@@ -17,7 +17,7 @@ name = "airc"
 path = "src/main.rs"
 
 [dependencies]
-air-ir = { package = "air-ir", path = "../ir", version = "0.4" }
+air-ir = { package = "air-ir", path = "../codegen/air", version = "0.4" }
 air-parser = { package = "air-parser", path = "../parser", version = "0.4" }
 air-pass = { package = "air-pass", path = "../pass", version = "0.1" }
 air-codegen-masm = { package = "air-codegen-masm", path = "../codegen/masm", version = "0.4" }

codegen/air/Cargo.toml

@@ -0,0 +1,20 @@
+[package]
+name = "air-ir"
+version = "0.4.0"
+description = "Intermediate representation for the AirScript language"
+authors = ["miden contributors"]
+readme = "README.md"
+license = "MIT"
+repository = "https://github.com/0xPolygonMiden/air-script"
+categories = ["compilers", "cryptography"]
+keywords = ["air", "stark", "zero-knowledge", "zkp"]
+rust-version.workspace = true
+edition.workspace = true
+
+[dependencies]
+air-parser = { package = "air-parser", path = "../../parser", version = "0.4" }
+air-pass = { package = "air-pass", path = "../../pass", version = "0.1" }
+mir = { package = "mir", path = "../../ir", version = "0.4" }
+anyhow = "1.0"
+miden-diagnostics = "0.1"
+thiserror = "1.0"

codegen/air/README.md

@@ -0,0 +1,35 @@
+# Intermediate Representation (IR)
+
+This crate contains the intermediate representation for AirScript, `AirIR`.
+
+The purpose of the `AirIR` is to provide a simple and accurate representation of an AIR that allows for optimization and translation to constraint evaluator code in a variety of target languages.
+
+## Generating the AirIR
+
+Generate an `AirIR` from an AirScript AST (the output of the AirScript parser) using the `new` method. The `new` method will return a new `AirIR` or an `Error` of type `SemanticError` if it encounters any errors while processing the AST.
+
+The `new` method will first iterate through the source sections that contain declarations to build a symbol table with constants, trace columns, public inputs, periodic columns and random values. It will return a `SemanticError` if it encounters a duplicate, incorrect, or missing declaration. Once the symbol table is built, the constraints and intermediate variables in the `boundary_constraints` and `integrity_constraints` sections of the AST are processed. Finally, `new` returns a Result containing the `AirIR` or a `SemanticError`.
+
+Example usage:
+
+```Rust
+// parse the source string to a Result containing the AST or an Error
+let ast = parse(source.as_str()).expect("Parsing failed");
+
+// process the AST to get a Result containing the AirIR or an Error
+let ir = AirIR::new(&ast)
+```
+
+## AirIR
+
+Although generation of an `AirIR` uses a symbol table while processing the source AST, the internal representation only consists of the following:
+
+- **Name** of the AIR definition represented by the `AirIR`.
+- **Segment Widths**, represented by a vector that contains the width of each trace segment (currently `main` and `auxiliary`).
+- **Constants**, represented by a vector that maps an identifier to a constant value.
+- **Public inputs**, represented by a vector that maps an identifier to a size for each public input that was declared. (Currently, public inputs can only be declared as fixed-size arrays.)
+- **Periodic columns**, represented by an ordered vector that contains each periodic column's repeating pattern (as a vector).
+- **Constraints**, represented by the combination of:
+  - a directed acyclic graph (DAG) without duplicate nodes.
+  - a vector of `ConstraintRoot` for each trace segment (e.g. main or auxiliary), where `ConstraintRoot` contains the node index in the graph where each of the constraint starts and the constraint domain which specifies the row(s) accessed by each of the constraints.
+  - contains both boundary and integrity constraints.

codegen/air/src/codegen.rs

@@ -0,0 +1,8 @@
+/// This trait should be implemented on types which handle generating code from AirScript IR
+pub trait CodeGenerator {
+    /// The type of the artifact produced by this codegen backend
+    type Output;
+
+    /// Generates code using this generator, consuming it in the process
+    fn generate(&self, ir: &crate::Air) -> anyhow::Result<Self::Output>;
+}

codegen/air/src/graph/mod.rs

@@ -0,0 +1,187 @@
+use std::collections::BTreeMap;
+
+use crate::ir::*;
+
+/// A unique identifier for a node in an [AlgebraicGraph]
+///
+/// The raw value of this identifier is an index in the `nodes` vector
+/// of the [AlgebraicGraph] struct.
+#[derive(Default, Debug, Clone, Copy, PartialEq, Eq)]
+pub struct NodeIndex(usize);
+impl core::ops::Add<usize> for NodeIndex {
+    type Output = NodeIndex;
+
+    fn add(self, rhs: usize) -> Self::Output {
+        Self(self.0 + rhs)
+    }
+}
+impl core::ops::Add<usize> for &NodeIndex {
+    type Output = NodeIndex;
+
+    fn add(self, rhs: usize) -> Self::Output {
+        NodeIndex(self.0 + rhs)
+    }
+}
+
+/// A node in the [AlgebraicGraph]
+#[derive(Debug, Clone)]
+pub struct Node {
+    /// The operation represented by this node
+    op: Operation,
+}
+impl Node {
+    /// Get the underlying [Operation] represented by this node
+    #[inline]
+    pub const fn op(&self) -> &Operation {
+        &self.op
+    }
+}
+
+/// The AlgebraicGraph is a directed acyclic graph used to represent integrity constraints. To
+/// store it compactly, it is represented as a vector of nodes where each node references other
+/// nodes by their index in the vector.
+///
+/// Within the graph, constraint expressions can overlap and share subgraphs, since new expressions
+/// reuse matching existing nodes when they are added, rather than creating new nodes.
+///
+/// - Leaf nodes (with no outgoing edges) are constants or references to trace cells (i.e. column 0
+///   in the current row or column 5 in the next row).
+/// - Tip nodes with no incoming edges (no parent nodes) always represent constraints, although they
+///   do not necessarily represent all constraints. There could be constraints which are also
+///   subgraphs of other constraints.
+#[derive(Default, Debug, Clone)]
+pub struct AlgebraicGraph {
+    /// All nodes in the graph.
+    nodes: Vec<Node>,
+}
+impl AlgebraicGraph {
+    /// Creates a new graph from a list of nodes.
+    pub const fn new(nodes: Vec<Node>) -> Self {
+        Self { nodes }
+    }
+
+    /// Returns the node with the specified index.
+    pub fn node(&self, index: &NodeIndex) -> &Node {
+        &self.nodes[index.0]
+    }
+
+    /// Returns the number of nodes in the graph.
+    pub fn num_nodes(&self) -> usize {
+        self.nodes.len()
+    }
+
+    /// Returns the degree of the subgraph which has the specified node as its tip.
+    pub fn degree(&self, index: &NodeIndex) -> IntegrityConstraintDegree {
+        let mut cycles = BTreeMap::default();
+        let base = self.accumulate_degree(&mut cycles, index);
+
+        if cycles.is_empty() {
+            IntegrityConstraintDegree::new(base)
+        } else {
+            IntegrityConstraintDegree::with_cycles(base, cycles.values().copied().collect())
+        }
+    }
+
+    /// TODO: docs
+    pub fn node_details(
+        &self,
+        index: &NodeIndex,
+        default_domain: ConstraintDomain,
+    ) -> Result<(TraceSegmentId, ConstraintDomain), ConstraintError> {
+        // recursively walk the subgraph and infer the trace segment and domain
+        match self.node(index).op() {
+            Operation::Value(value) => match value {
+                Value::Constant(_) => Ok((DEFAULT_SEGMENT, default_domain)),
+                Value::PeriodicColumn(_) => {
+                    assert!(
+                        !default_domain.is_boundary(),
+                        "unexpected access to periodic column in boundary constraint"
+                    );
+                    // the default domain for [IntegrityConstraints] is `EveryRow`
+                    Ok((DEFAULT_SEGMENT, ConstraintDomain::EveryRow))
+                }
+                Value::PublicInput(_) => {
+                    assert!(
+                        !default_domain.is_integrity(),
+                        "unexpected access to public input in integrity constraint"
+                    );
+                    Ok((DEFAULT_SEGMENT, default_domain))
+                }
+                Value::RandomValue(_) => Ok((AUX_SEGMENT, default_domain)),
+                Value::TraceAccess(trace_access) => {
+                    let domain = if default_domain.is_boundary() {
+                        assert_eq!(
+                            trace_access.row_offset, 0,
+                            "unexpected trace offset in boundary constraint"
+                        );
+                        default_domain
+                    } else {
+                        ConstraintDomain::from_offset(trace_access.row_offset)
+                    };
+
+                    Ok((trace_access.segment, domain))
+                }
+            },
+            Operation::Add(lhs, rhs) | Operation::Sub(lhs, rhs) | Operation::Mul(lhs, rhs) => {
+                let (lhs_segment, lhs_domain) = self.node_details(lhs, default_domain)?;
+                let (rhs_segment, rhs_domain) = self.node_details(rhs, default_domain)?;
+
+                let trace_segment = lhs_segment.max(rhs_segment);
+                let domain = lhs_domain.merge(rhs_domain)?;
+
+                Ok((trace_segment, domain))
+            }
+        }
+    }
+
+    /// Insert the operation and return its node index. If an identical node already exists, return
+    /// that index instead.
+    pub(crate) fn insert_node(&mut self, op: Operation) -> NodeIndex {
+        self.nodes.iter().position(|n| *n.op() == op).map_or_else(
+            || {
+                // create a new node.
+                let index = self.nodes.len();
+                self.nodes.push(Node { op });
+                NodeIndex(index)
+            },
+            |index| {
+                // return the existing node's index.
+                NodeIndex(index)
+            },
+        )
+    }
+
+    /// Recursively accumulates the base degree and the cycle lengths of the periodic columns.
+    fn accumulate_degree(
+        &self,
+        cycles: &mut BTreeMap<QualifiedIdentifier, usize>,
+        index: &NodeIndex,
+    ) -> usize {
+        // recursively walk the subgraph and compute the degree from the operation and child nodes
+        match self.node(index).op() {
+            Operation::Value(value) => match value {
+                Value::Constant(_) | Value::RandomValue(_) | Value::PublicInput(_) => 0,
+                Value::TraceAccess(_) => 1,
+                Value::PeriodicColumn(pc) => {
+                    cycles.insert(pc.name, pc.cycle);
+                    0
+                }
+            },
+            Operation::Add(lhs, rhs) => {
+                let lhs_base = self.accumulate_degree(cycles, lhs);
+                let rhs_base = self.accumulate_degree(cycles, rhs);
+                lhs_base.max(rhs_base)
+            }
+            Operation::Sub(lhs, rhs) => {
+                let lhs_base = self.accumulate_degree(cycles, lhs);
+                let rhs_base = self.accumulate_degree(cycles, rhs);
+                lhs_base.max(rhs_base)
+            }
+            Operation::Mul(lhs, rhs) => {
+                let lhs_base = self.accumulate_degree(cycles, lhs);
+                let rhs_base = self.accumulate_degree(cycles, rhs);
+                lhs_base + rhs_base
+            }
+        }
+    }
+}