Replace indices: HashMap with indices: Fn(T)->usize

This commit is contained in:
wborgeaud 2021-07-02 14:37:07 +02:00
parent d93cf693ba
commit b6554ba2ec
2 changed files with 15 additions and 13 deletions

View File

@ -359,7 +359,13 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
fn sigma_vecs(&self, k_is: &[F], subgroup: &[F]) -> Vec<PolynomialValues<F>> { fn sigma_vecs(&self, k_is: &[F], subgroup: &[F]) -> Vec<PolynomialValues<F>> {
let degree = self.gate_instances.len(); let degree = self.gate_instances.len();
let degree_log = log2_strict(degree); let degree_log = log2_strict(degree);
let mut target_partition = TargetPartition::default(); let mut target_partition = TargetPartition::new(|t| match t {
Target::Wire(Wire { gate, input }) => gate * self.config.num_routed_wires + input,
Target::PublicInput { index } => degree * self.config.num_routed_wires + index,
Target::VirtualTarget { index } => {
degree * self.config.num_routed_wires + self.public_input_index + index
}
});
for gate in 0..degree { for gate in 0..degree {
for input in 0..self.config.num_routed_wires { for input in 0..self.config.num_routed_wires {

View File

@ -18,21 +18,18 @@ pub struct ForestNode<T: Debug + Copy + Eq + PartialEq> {
} }
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub struct TargetPartition<T: Debug + Copy + Eq + PartialEq + Hash> { pub struct TargetPartition<T: Debug + Copy + Eq + PartialEq + Hash, F: Fn(T) -> usize> {
forest: Vec<ForestNode<T>>, forest: Vec<ForestNode<T>>,
indices: HashMap<T, usize>, indices: F,
} }
impl<T: Debug + Copy + Eq + PartialEq + Hash> Default for TargetPartition<T> { impl<T: Debug + Copy + Eq + PartialEq + Hash, F: Fn(T) -> usize> TargetPartition<T, F> {
fn default() -> Self { pub fn new(f: F) -> Self {
Self { Self {
forest: Vec::new(), forest: Vec::new(),
indices: Default::default(), indices: f,
} }
} }
}
impl<T: Debug + Copy + Eq + PartialEq + Hash> TargetPartition<T> {
/// Add a new partition with a single member. /// Add a new partition with a single member.
pub fn add(&mut self, t: T) { pub fn add(&mut self, t: T) {
let index = self.forest.len(); let index = self.forest.len();
@ -42,7 +39,6 @@ impl<T: Debug + Copy + Eq + PartialEq + Hash> TargetPartition<T> {
size: 1, size: 1,
index, index,
}); });
self.indices.insert(t, index);
} }
/// Path halving /// Path halving
@ -58,8 +54,8 @@ impl<T: Debug + Copy + Eq + PartialEq + Hash> TargetPartition<T> {
/// Merge the two partitions containing the two given targets. Does nothing if the targets are /// Merge the two partitions containing the two given targets. Does nothing if the targets are
/// already members of the same partition. /// already members of the same partition.
pub fn merge(&mut self, tx: T, ty: T) { pub fn merge(&mut self, tx: T, ty: T) {
let index_x = self.indices[&tx]; let index_x = (self.indices)(tx);
let index_y = self.indices[&ty]; let index_y = (self.indices)(ty);
let mut x = self.forest[index_x]; let mut x = self.forest[index_x];
let mut y = self.forest[index_y]; let mut y = self.forest[index_y];
@ -81,7 +77,7 @@ impl<T: Debug + Copy + Eq + PartialEq + Hash> TargetPartition<T> {
self.forest[index_y] = y; self.forest[index_y] = y;
} }
} }
impl TargetPartition<Target> { impl<F: Fn(Target) -> usize> TargetPartition<Target, F> {
pub fn wire_partitions(&mut self) -> WirePartitions { pub fn wire_partitions(&mut self) -> WirePartitions {
let mut partition = HashMap::<_, Vec<_>>::new(); let mut partition = HashMap::<_, Vec<_>>::new();
let nodes = self.forest.clone(); let nodes = self.forest.clone();