From 46b5ddcc687792e71cb1c426064bbeba04263d12 Mon Sep 17 00:00:00 2001
From: therealyingtong <yingtong@z.cash>
Date: Fri, 15 Jan 2021 13:09:26 +0800
Subject: [PATCH] Compression subregion_digest assignments

---
 src/gadget/sha256/table16/compression.rs      |  12 +-
 .../table16/compression/subregion_digest.rs   | 122 ++++++++++++++++++
 2 files changed, 132 insertions(+), 2 deletions(-)
 create mode 100644 src/gadget/sha256/table16/compression/subregion_digest.rs

diff --git a/src/gadget/sha256/table16/compression.rs b/src/gadget/sha256/table16/compression.rs
index 7547bdd1..58ad8181 100644
--- a/src/gadget/sha256/table16/compression.rs
+++ b/src/gadget/sha256/table16/compression.rs
@@ -11,7 +11,7 @@ use crate::{
 
 mod compression_gates;
 mod compression_util;
-// mod subregion_digest;
+mod subregion_digest;
 mod subregion_initial;
 mod subregion_main;
 
@@ -738,7 +738,15 @@ impl Compression {
         state: State,
     ) -> Result<[BlockWord; DIGEST_SIZE], Error> {
         let mut digest = [BlockWord::new(0); DIGEST_SIZE];
-        todo!()
+        layouter.assign_region(
+            || "digest",
+            |mut region| {
+                digest = self.assign_digest(&mut region, state.clone())?;
+
+                Ok(())
+            },
+        )?;
+        Ok(digest)
     }
 
     pub(super) fn empty_configure<F: FieldExt>(
diff --git a/src/gadget/sha256/table16/compression/subregion_digest.rs b/src/gadget/sha256/table16/compression/subregion_digest.rs
new file mode 100644
index 00000000..dfcb3040
--- /dev/null
+++ b/src/gadget/sha256/table16/compression/subregion_digest.rs
@@ -0,0 +1,122 @@
+use super::super::{super::DIGEST_SIZE, BlockWord, CellValue16, Table16Assignment, Table16Chip};
+use super::{compression_util::*, Compression, State};
+use crate::{
+    arithmetic::FieldExt,
+    circuit::Region,
+    plonk::{Advice, Column, Error},
+};
+
+impl Compression {
+    pub fn assign_digest<F: FieldExt>(
+        &self,
+        region: &mut Region<'_, Table16Chip<F>>,
+        state: State,
+    ) -> Result<[BlockWord; DIGEST_SIZE], Error> {
+        let a_3 = self.extras[0];
+        let a_4 = self.extras[1];
+        let a_5 = self.message_schedule;
+        let a_6 = self.extras[2];
+        let a_7 = self.extras[3];
+        let a_8 = self.extras[4];
+
+        let (a, b, c, d, e, f, g, h) = match_state(state);
+
+        let abcd_row = 0;
+        region.assign_fixed(|| "s_digest", self.s_digest, abcd_row, || Ok(F::one()))?;
+        let efgh_row = abcd_row + 2;
+        region.assign_fixed(|| "s_digest", self.s_digest, efgh_row, || Ok(F::one()))?;
+
+        // Assign digest for A, B, C, D
+        self.assign_and_constrain(
+            region,
+            || "a_lo",
+            a_3,
+            abcd_row,
+            &a.dense_halves.0.into(),
+            &self.perm,
+        )?;
+        self.assign_and_constrain(
+            region,
+            || "a_hi",
+            a_4,
+            abcd_row,
+            &a.dense_halves.1.into(),
+            &self.perm,
+        )?;
+        let a = a.dense_halves.0.value.unwrap() as u32
+            + (1 << 16) * (a.dense_halves.1.value.unwrap() as u32);
+        region.assign_advice(|| "a", a_5, abcd_row, || Ok(F::from_u64(a as u64)))?;
+
+        let b = self.assign_digest_word(region, abcd_row, a_6, a_7, a_8, b.dense_halves)?;
+        let c = self.assign_digest_word(region, abcd_row + 1, a_3, a_4, a_5, c.dense_halves)?;
+        let d = self.assign_digest_word(region, abcd_row + 1, a_6, a_7, a_8, d.dense_halves)?;
+
+        // Assign digest for E, F, G, H
+        self.assign_and_constrain(
+            region,
+            || "e_lo",
+            a_3,
+            efgh_row,
+            &e.dense_halves.0.into(),
+            &self.perm,
+        )?;
+        self.assign_and_constrain(
+            region,
+            || "e_hi",
+            a_4,
+            efgh_row,
+            &e.dense_halves.1.into(),
+            &self.perm,
+        )?;
+        let e = e.dense_halves.0.value.unwrap() as u32
+            + (1 << 16) * (e.dense_halves.1.value.unwrap() as u32);
+        region.assign_advice(|| "e", a_5, efgh_row, || Ok(F::from_u64(e as u64)))?;
+
+        let f = self.assign_digest_word(region, efgh_row, a_6, a_7, a_8, f.dense_halves)?;
+        let g = self.assign_digest_word(region, efgh_row + 1, a_3, a_4, a_5, g.dense_halves)?;
+        let h = self.assign_digest_word(region, efgh_row + 1, a_6, a_7, a_8, h.dense_halves)?;
+
+        Ok([
+            BlockWord::new(a),
+            BlockWord::new(b),
+            BlockWord::new(c),
+            BlockWord::new(d),
+            BlockWord::new(e),
+            BlockWord::new(f),
+            BlockWord::new(g),
+            BlockWord::new(h),
+        ])
+    }
+
+    fn assign_digest_word<F: FieldExt>(
+        &self,
+        region: &mut Region<'_, Table16Chip<F>>,
+        row: usize,
+        lo_col: Column<Advice>,
+        hi_col: Column<Advice>,
+        word_col: Column<Advice>,
+        dense_halves: (CellValue16, CellValue16),
+    ) -> Result<u32, Error> {
+        self.assign_and_constrain(
+            region,
+            || "lo",
+            lo_col,
+            row,
+            &dense_halves.0.into(),
+            &self.perm,
+        )?;
+        self.assign_and_constrain(
+            region,
+            || "hi",
+            hi_col,
+            row,
+            &dense_halves.1.into(),
+            &self.perm,
+        )?;
+        let val = dense_halves.0.value.unwrap() as u32
+            + (1 << 16) * (dense_halves.1.value.unwrap() as u32);
+        region.assign_advice(|| "word", word_col, row, || Ok(F::from_u64(val as u64)))?;
+
+        Ok(val)
+    }
+}