feat!: Introduce WitnessMap data structure to avoid leaking internal …

…structure (#252)

* feat!: Introduce WitnessMap data structure to avoid leaking internal structure

* avoid fully qualified name

* Implement IntoIterator for WitnessMap

* Make WitnessMap its own module

* move thiserror to the workspace

* Implement try_from instead of from where there was unwrapping

* clippy

---------

Co-authored-by: Tom French <tom@tomfren.ch>

Cargo.toml

@@ -16,4 +16,6 @@ stdlib = { package = "acvm_stdlib", version = "0.11.0", path = "stdlib", default
 num-bigint = "0.4"
 num-traits = "0.2"
 
+thiserror = "1.0.21"
+
 serde = { version = "1.0.136", features = ["derive"] }

acir/Cargo.toml

@@ -12,6 +12,7 @@ rust-version.workspace = true
 [dependencies]
 acir_field.workspace = true
 serde.workspace = true
+thiserror.workspace = true
 
 rmp-serde = "1.1.0"
 flate2 = "1.0.24"

acir/src/native_types/mod.rs

@@ -1,5 +1,7 @@
 mod expression;
 mod witness;
+mod witness_map;
 
 pub use expression::Expression;
 pub use witness::Witness;
+pub use witness_map::WitnessMap;

acir/src/native_types/witness.rs

@@ -1,9 +1,3 @@
-use std::io::Read;
-
-use flate2::{
-    bufread::{DeflateDecoder, DeflateEncoder},
-    Compression,
-};
 use serde::{Deserialize, Serialize};
 
 // Witness might be a misnomer. This is an index that represents the position a witness will take
@@ -27,23 +21,10 @@ impl Witness {
     pub const fn can_defer_constraint(&self) -> bool {
         true
     }
+}
 
-    pub fn to_bytes(
-        witnesses: &std::collections::BTreeMap<Witness, acir_field::FieldElement>,
-    ) -> Vec<u8> {
-        let buf = rmp_serde::to_vec(witnesses).unwrap();
-        let mut deflater = DeflateEncoder::new(buf.as_slice(), Compression::best());
-        let mut buf_c = Vec::new();
-        deflater.read_to_end(&mut buf_c).unwrap();
-        buf_c
-    }
-
-    pub fn from_bytes(
-        bytes: &[u8],
-    ) -> std::collections::BTreeMap<Witness, acir_field::FieldElement> {
-        let mut deflater = DeflateDecoder::new(bytes);
-        let mut buf_d = Vec::new();
-        deflater.read_to_end(&mut buf_d).unwrap();
-        rmp_serde::from_slice(buf_d.as_slice()).unwrap()
+impl From<u32> for Witness {
+    fn from(value: u32) -> Self {
+        Self(value)
     }
 }

acir/src/native_types/witness_map.rs

@@ -0,0 +1,105 @@
+use std::{
+    collections::{btree_map, BTreeMap},
+    io::Read,
+    ops::Index,
+};
+
+use acir_field::FieldElement;
+use flate2::{
+    bufread::{DeflateDecoder, DeflateEncoder},
+    Compression,
+};
+use serde::{Deserialize, Serialize};
+use thiserror::Error;
+
+use crate::native_types::Witness;
+
+#[derive(Debug, Error)]
+pub enum WitnessMapError {
+    #[error(transparent)]
+    MsgpackEncode(#[from] rmp_serde::encode::Error),
+
+    #[error(transparent)]
+    MsgpackDecode(#[from] rmp_serde::decode::Error),
+
+    #[error(transparent)]
+    Deflate(#[from] std::io::Error),
+}
+
+/// A map from the witnesses in a constraint system to the field element values
+#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Default, Serialize, Deserialize)]
+pub struct WitnessMap(BTreeMap<Witness, FieldElement>);
+
+impl WitnessMap {
+    pub fn new() -> Self {
+        Self(BTreeMap::new())
+    }
+    pub fn get(&self, witness: &Witness) -> Option<&FieldElement> {
+        self.0.get(witness)
+    }
+    pub fn get_index(&self, index: u32) -> Option<&FieldElement> {
+        self.0.get(&index.into())
+    }
+    pub fn contains_key(&self, key: &Witness) -> bool {
+        self.0.contains_key(key)
+    }
+    pub fn insert(&mut self, key: Witness, value: FieldElement) -> Option<FieldElement> {
+        self.0.insert(key, value)
+    }
+}
+
+impl Index<&Witness> for WitnessMap {
+    type Output = FieldElement;
+
+    fn index(&self, index: &Witness) -> &Self::Output {
+        &self.0[index]
+    }
+}
+
+pub struct IntoIter(btree_map::IntoIter<Witness, FieldElement>);
+
+impl Iterator for IntoIter {
+    type Item = (Witness, FieldElement);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.0.next()
+    }
+}
+
+impl IntoIterator for WitnessMap {
+    type Item = (Witness, FieldElement);
+    type IntoIter = IntoIter;
+
+    fn into_iter(self) -> Self::IntoIter {
+        IntoIter(self.0.into_iter())
+    }
+}
+
+impl From<BTreeMap<Witness, FieldElement>> for WitnessMap {
+    fn from(value: BTreeMap<Witness, FieldElement>) -> Self {
+        Self(value)
+    }
+}
+
+impl TryFrom<WitnessMap> for Vec<u8> {
+    type Error = WitnessMapError;
+
+    fn try_from(val: WitnessMap) -> Result<Self, Self::Error> {
+        let buf = rmp_serde::to_vec(&val)?;
+        let mut deflater = DeflateEncoder::new(buf.as_slice(), Compression::best());
+        let mut buf_c = Vec::new();
+        deflater.read_to_end(&mut buf_c)?;
+        Ok(buf_c)
+    }
+}
+
+impl TryFrom<&[u8]> for WitnessMap {
+    type Error = WitnessMapError;
+
+    fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
+        let mut deflater = DeflateDecoder::new(bytes);
+        let mut buf_d = Vec::new();
+        deflater.read_to_end(&mut buf_d)?;
+        Ok(Self(rmp_serde::from_slice(buf_d.as_slice())?))
+    }
+}

acvm/Cargo.toml

@@ -12,6 +12,7 @@ rust-version.workspace = true
 [dependencies]
 num-bigint.workspace = true
 num-traits.workspace = true
+thiserror.workspace = true
 
 acir.workspace = true
 stdlib.workspace = true
@@ -28,7 +29,6 @@ k256 = { version = "0.7.2", features = [
     "arithmetic",
 ] }
 indexmap = "1.7.0"
-thiserror = "1.0.21"
 
 [features]
 default = ["bn254"]