From 7b8c4fce219ba06e0f168b92b6304bdb359811be Mon Sep 17 00:00:00 2001
From: Georg Wiese <georgwiese@gmail.com>
Date: Mon, 17 Jul 2023 17:47:40 +0000
Subject: [PATCH] Fix using constants in identities #401

---
 compiler/tests/pil.rs                          |  7 +++++++
 halo2/src/circuit_builder.rs                   |  5 ++++-
 halo2/src/circuit_data.rs                      |  8 +++++++-
 pil_analyzer/src/json_exporter/mod.rs          | 11 ++++-------
 test_data/asm/secondary_block_machine_add2.asm |  5 ++++-
 test_data/pil/constant_in_identity.pil         | 16 ++++++++++++++++
 6 files changed, 42 insertions(+), 10 deletions(-)
 create mode 100644 test_data/pil/constant_in_identity.pil

diff --git a/compiler/tests/pil.rs b/compiler/tests/pil.rs
index d12a0c672..d5b5678c9 100644
--- a/compiler/tests/pil.rs
+++ b/compiler/tests/pil.rs
@@ -36,6 +36,13 @@ fn test_fibonacci() {
     halo2_proof(f, Default::default());
 }
 
+#[test]
+fn test_constant_in_identity() {
+    let f = "constant_in_identity.pil";
+    verify_pil(f, None);
+    halo2_proof(f, Default::default());
+}
+
 #[test]
 fn test_fibonacci_macro() {
     let f = "fib_macro.pil";
diff --git a/halo2/src/circuit_builder.rs b/halo2/src/circuit_builder.rs
index a362d64c2..ad34be19c 100644
--- a/halo2/src/circuit_builder.rs
+++ b/halo2/src/circuit_builder.rs
@@ -36,7 +36,7 @@ pub(crate) fn analyzed_to_circuit<T: FieldElement>(
 
     let query = |column, rotation| Expr::Var(PlonkVar::Query(ColumnQuery { column, rotation }));
 
-    let mut cd = CircuitData::from(fixed, witness);
+    let mut cd = CircuitData::from(fixed, witness, &analyzed.constants);
 
     // append two fixed columns:
     // - one that enables constraints that do not have rotations (__enable_cur) in the actual circuit
@@ -253,6 +253,9 @@ fn expression_2_expr<T: FieldElement>(cd: &CircuitData<T>, expr: &Expression<T>)
                 _ => unimplemented!("{:?}", expr),
             }
         }
+        Expression::Constant(constant_name) => {
+            Expr::Const(cd.constants[constant_name].to_arbitrary_integer())
+        }
 
         _ => unimplemented!("{:?}", expr),
     }
diff --git a/halo2/src/circuit_data.rs b/halo2/src/circuit_data.rs
index 27afd74a4..68ed0c669 100644
--- a/halo2/src/circuit_data.rs
+++ b/halo2/src/circuit_data.rs
@@ -10,10 +10,15 @@ pub(crate) struct CircuitData<'a, T> {
     pub(crate) fixed: Vec<(&'a str, Vec<T>)>,
     pub(crate) witness: Vec<(&'a str, Vec<T>)>,
     columns: HashMap<String, Column>,
+    pub(crate) constants: &'a HashMap<String, T>,
 }
 
 impl<'a, T: FieldElement> CircuitData<'a, T> {
-    pub fn from(fixed: Vec<(&'a str, Vec<T>)>, witness: Vec<(&'a str, Vec<T>)>) -> Self {
+    pub fn from(
+        fixed: Vec<(&'a str, Vec<T>)>,
+        witness: Vec<(&'a str, Vec<T>)>,
+        constants: &'a HashMap<String, T>,
+    ) -> Self {
         if !fixed.is_empty() && !witness.is_empty() {
             assert_eq!(
                 fixed.get(0).unwrap().1.len(),
@@ -47,6 +52,7 @@ impl<'a, T: FieldElement> CircuitData<'a, T> {
             fixed,
             witness,
             columns,
+            constants,
         }
     }
 
diff --git a/pil_analyzer/src/json_exporter/mod.rs b/pil_analyzer/src/json_exporter/mod.rs
index 1e8f6beee..4d99363fa 100644
--- a/pil_analyzer/src/json_exporter/mod.rs
+++ b/pil_analyzer/src/json_exporter/mod.rs
@@ -217,14 +217,11 @@ impl<'a, T: FieldElement> Exporter<'a, T> {
     fn expression_to_json(&self, expr: &Expression<T>) -> (u32, JsonValue, Vec<u64>) {
         match expr {
             Expression::Constant(name) => (
-                1,
+                0,
                 object! {
-                    // TODO I think "const" is for constant poly, not a constant value.
-                    op: "const",
-                    deg: 1,
-                    // TODO is it declarations or constants?
-                    id: self.analyzed.definitions[name].0.id,
-                    next: false
+                    op: "number",
+                    deg: 0,
+                    value: format!("{}", self.analyzed.constants[name]),
                 },
                 Vec::new(),
             ),
diff --git a/test_data/asm/secondary_block_machine_add2.asm b/test_data/asm/secondary_block_machine_add2.asm
index bff2afb51..0bc3e0056 100644
--- a/test_data/asm/secondary_block_machine_add2.asm
+++ b/test_data/asm/secondary_block_machine_add2.asm
@@ -19,9 +19,12 @@ machine Main {
         // Because constraints are not cyclic, we need to explicitly constrain the first state
         first_step * (add_two_state - add_two_input) = 0;
 
+        // Add %offset in a single step of computation
+        constant %offset = 1;
+
         // If RESET is true, constrain the next state to be equal to the input
         // if RESET is false, increment the current state
-        add_two_state' = (1 - add_two_RESET) * (add_two_state + 1) + add_two_RESET * add_two_input';
+        add_two_state' = (1 - add_two_RESET) * (add_two_state + %offset) + add_two_RESET * add_two_input';
 
         // If RESET is true, the next input is unconstrained
         // If RESET is false, the next input is equal to the current input
diff --git a/test_data/pil/constant_in_identity.pil b/test_data/pil/constant_in_identity.pil
new file mode 100644
index 000000000..ab0532606
--- /dev/null
+++ b/test_data/pil/constant_in_identity.pil
@@ -0,0 +1,16 @@
+constant %N = 16;
+
+namespace ConstantInIdentity(%N);
+    col fixed ISLAST(i) { match i {
+        %N - 1 => 1,
+        _ => 0,
+    } };
+    col witness x;
+
+    constant %offset = 5;
+
+    ISLAST * (x' - 1) = 0;
+
+    (1-ISLAST) * (x' - x - %offset) = 0;
+
+    public out = x(%N-1);