Merge pull request #238 from varunthakore/generic-log

Implement generic log

Author: katat

src/backends/kimchi/mod.rs

@@ -44,6 +44,8 @@ pub const NUM_REGISTERS: usize = kimchi::circuits::wires::COLUMNS;
 
 use super::{Backend, BackendField, BackendVar};
 
+use crate::mast::Mast;
+
 impl BackendField for VestaField {
     fn to_circ_field(&self) -> circ_fields::FieldV {
         let mut opt = CircOpt::default();
@@ -416,10 +418,11 @@ impl Backend for KimchiVesta {
         self.compute_val(env, val, var.index)
     }
 
-    fn generate_witness(
+    fn generate_witness<B: Backend>(
         &self,
         witness_env: &mut WitnessEnv<VestaField>,
         sources: &Sources,
+        _typed: &Mast<B>,
     ) -> Result<GeneratedWitness> {
         if !self.finalized {
             unreachable!("the circuit must be finalized before generating a witness");

src/backends/mod.rs

@@ -22,6 +22,8 @@ use self::{
     r1cs::{R1csBls12381Field, R1csBn254Field, R1CS},
 };
 
+use crate::mast::Mast;
+
 pub mod kimchi;
 pub mod r1cs;
 
@@ -411,10 +413,11 @@ pub trait Backend: Clone {
     ) -> Result<()>;
 
     /// Generate the witness for a backend.
-    fn generate_witness(
+    fn generate_witness<B: Backend>(
         &self,
         witness_env: &mut WitnessEnv<Self::Field>,
         sources: &Sources,
+        typed: &Mast<B>,
     ) -> Result<Self::GeneratedWitness>;
 
     /// Generate the asm for a backend.

src/backends/r1cs/mod.rs

@@ -17,6 +17,9 @@ use crate::circuit_writer::VarInfo;
 use crate::compiler::Sources;
 use crate::constants::Span;
 use crate::error::{Error, ErrorKind, Result};
+use crate::mast::Mast;
+use crate::parser::types::{ModulePath, TyKind};
+use crate::type_checker::FullyQualified;
 use crate::var::ConstOrCell;
 use crate::{circuit_writer::DebugInfo, var::Value};
 
@@ -485,10 +488,11 @@ where
 
     /// Generate the witnesses
     /// This process should check if the constraints are satisfied.
-    fn generate_witness(
+    fn generate_witness<B: Backend>(
         &self,
         witness_env: &mut crate::witness::WitnessEnv<F>,
         sources: &Sources,
+        typed: &Mast<B>,
     ) -> crate::error::Result<Self::GeneratedWitness> {
         assert!(self.finalized, "the circuit is not finalized yet!");
 
@@ -520,15 +524,69 @@ where
             let (line, _, line_str) = crate::utils::find_exact_line(source, *span);
             let line_str = line_str.trim_start();
             let dbg_msg = format!("[{filename}:{line}] `{line_str}` -> ");
-            for cvar in var_info.var.iter() {
-                match cvar {
+
+            match &var_info.typ {
+                // Field
+                Some(TyKind::Field { .. }) => match &var_info.var[0] {
                     ConstOrCell::Const(cst) => {
                         println!("{dbg_msg}{}", cst.pretty());
                     }
                     ConstOrCell::Cell(cell) => {
                         let val = cell.evaluate(&witness);
                         println!("{dbg_msg}{}", val.pretty());
                     }
+                },
+
+                // Bool
+                Some(TyKind::Bool) => match &var_info.var[0] {
+                    ConstOrCell::Const(cst) => {
+                        let val = *cst == F::one();
+                        println!("{dbg_msg}{}", val);
+                    }
+                    ConstOrCell::Cell(cell) => {
+                        let val = cell.evaluate(&witness) == F::one();
+                        println!("{dbg_msg}{}", val);
+                    }
+                },
+
+                // Array
+                Some(TyKind::Array(b, s)) => {
+                    let (output, remaining) =
+                        log_array_type(&var_info.var.cvars, b, *s, &witness, typed, span);
+                    assert!(remaining.is_empty());
+                    println!("{dbg_msg}{}", output);
+                }
+
+                // Custom types
+                Some(TyKind::Custom {
+                    module,
+                    name: struct_name,
+                }) => {
+                    let mut string_vec = Vec::new();
+                    let (output, remaining) = log_custom_type(
+                        module,
+                        struct_name,
+                        typed,
+                        &var_info.var.cvars,
+                        &witness,
+                        span,
+                        &mut string_vec,
+                    );
+                    assert!(remaining.is_empty());
+                    println!("{dbg_msg}{}{}", struct_name, output);
+                }
+
+                // GenericSizedArray
+                Some(TyKind::GenericSizedArray(_, _)) => {
+                    unreachable!("GenericSizedArray should be monomorphized")
+                }
+
+                None => {
+                    return Err(Error::new(
+                        "log",
+                        ErrorKind::UnexpectedError("No type info for logging"),
+                        *span,
+                    ))
                 }
             }
         }
@@ -704,6 +762,173 @@ where
     }
 }
 
+fn log_custom_type<F: BackendField, B: Backend>(
+    module: &ModulePath,
+    struct_name: &String,
+    typed: &Mast<B>,
+    var_info_var: &[ConstOrCell<F, LinearCombination<F>>],
+    witness: &[F],
+    span: &Span,
+    string_vec: &mut Vec<String>,
+) -> (String, Vec<ConstOrCell<F, LinearCombination<F>>>) {
+    let qualified = FullyQualified::new(module, struct_name);
+    let struct_info = typed
+        .struct_info(&qualified)
+        .ok_or(
+            typed
+                .0
+                .error(ErrorKind::UnexpectedError("struct not found"), *span),
+        )
+        .unwrap();
+
+    let mut remaining = var_info_var.to_vec();
+
+    for (field_name, field_typ) in &struct_info.fields {
+        let len = typed.size_of(field_typ);
+        match field_typ {
+            TyKind::Field { .. } => match &remaining[0] {
+                ConstOrCell::Const(cst) => {
+                    string_vec.push(format!("{field_name}: {}", cst.pretty()));
+                    remaining = remaining[len..].to_vec();
+                }
+                ConstOrCell::Cell(cell) => {
+                    let val = cell.evaluate(witness);
+                    string_vec.push(format!("{field_name}: {}", val.pretty()));
+                    remaining = remaining[len..].to_vec();
+                }
+            },
+
+            TyKind::Bool => match &remaining[0] {
+                ConstOrCell::Const(cst) => {
+                    let val = *cst == F::one();
+                    string_vec.push(format!("{field_name}: {}", val));
+                    remaining = remaining[len..].to_vec();
+                }
+                ConstOrCell::Cell(cell) => {
+                    let val = cell.evaluate(witness) == F::one();
+                    string_vec.push(format!("{field_name}: {}", val));
+                    remaining = remaining[len..].to_vec();
+                }
+            },
+
+            TyKind::Array(b, s) => {
+                let (output, new_remaining) =
+                    log_array_type(&remaining, b, *s, witness, typed, span);
+                string_vec.push(format!("{field_name}: {}", output));
+                remaining = new_remaining;
+            }
+
+            TyKind::Custom {
+                module,
+                name: struct_name,
+            } => {
+                let mut custom_string_vec = Vec::new();
+                let (output, new_remaining) = log_custom_type(
+                    module,
+                    struct_name,
+                    typed,
+                    &remaining,
+                    witness,
+                    span,
+                    &mut custom_string_vec,
+                );
+                string_vec.push(format!("{}: {}{}", field_name, struct_name, output));
+                remaining = new_remaining;
+            }
+
+            TyKind::GenericSizedArray(_, _) => {
+                unreachable!("GenericSizedArray should be monomorphized")
+            }
+        }
+    }
+
+    (format!("{{ {} }}", string_vec.join(", ")), remaining)
+}
+
+fn log_array_type<F: BackendField, B: Backend>(
+    var_info_var: &[ConstOrCell<F, LinearCombination<F>>],
+    base_type: &TyKind,
+    size: u32,
+    witness: &[F],
+    typed: &Mast<B>,
+    span: &Span,
+) -> (String, Vec<ConstOrCell<F, LinearCombination<F>>>) {
+    match base_type {
+        TyKind::Field { .. } => {
+            let values: Vec<String> = var_info_var
+                .iter()
+                .take(size as usize)
+                .map(|cvar| match cvar {
+                    ConstOrCell::Const(cst) => cst.pretty(),
+                    ConstOrCell::Cell(cell) => cell.evaluate(witness).pretty(),
+                })
+                .collect();
+
+            let remaining = var_info_var[size as usize..].to_vec();
+            (format!("[{}]", values.join(", ")), remaining)
+        }
+
+        TyKind::Bool => {
+            let values: Vec<String> = var_info_var
+                .iter()
+                .take(size as usize)
+                .map(|cvar| match cvar {
+                    ConstOrCell::Const(cst) => {
+                        let val = *cst == F::one();
+                        val.to_string()
+                    }
+                    ConstOrCell::Cell(cell) => {
+                        let val = cell.evaluate(witness) == F::one();
+                        val.to_string()
+                    }
+                })
+                .collect();
+
+            let remaining = var_info_var[size as usize..].to_vec();
+            (format!("[{}]", values.join(", ")), remaining)
+        }
+
+        TyKind::Array(inner_type, inner_size) => {
+            let mut nested_result = Vec::new();
+            let mut remaining = var_info_var.to_vec();
+            for _ in 0..size {
+                let (chunk_result, new_remaining) =
+                    log_array_type(&remaining, inner_type, *inner_size, witness, typed, span);
+                nested_result.push(chunk_result);
+                remaining = new_remaining;
+            }
+            (format!("[{}]", nested_result.join(", ")), remaining)
+        }
+
+        TyKind::Custom {
+            module,
+            name: struct_name,
+        } => {
+            let mut nested_result = Vec::new();
+            let mut remaining = var_info_var.to_vec();
+            for _ in 0..size {
+                let mut string_vec = Vec::new();
+                let (output, new_remaining) = log_custom_type(
+                    module,
+                    struct_name,
+                    typed,
+                    &remaining,
+                    witness,
+                    span,
+                    &mut string_vec,
+                );
+                nested_result.push(format!("{}{}", struct_name, output));
+                remaining = new_remaining;
+            }
+            (format!("[{}]", nested_result.join(", ")), remaining)
+        }
+
+        TyKind::GenericSizedArray(_, _) => {
+            unreachable!("GenericSizedArray should be monomorphized")
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::{

src/circuit_writer/mod.rs

@@ -232,7 +232,8 @@ impl<B: Backend> CircuitWriter<B> {
         witness_env: &mut WitnessEnv<B::Field>,
         sources: &Sources,
     ) -> Result<B::GeneratedWitness> {
-        self.backend.generate_witness(witness_env, sources)
+        self.backend
+            .generate_witness(witness_env, sources, &self.typed)
     }
 
     fn handle_arg(