From b600142cd454b95eba403fa1f86f582ff8688c79 Mon Sep 17 00:00:00 2001
From: Robin Salen <30937548+Nashtare@users.noreply.github.com>
Date: Sat, 10 Feb 2024 15:48:52 -0500
Subject: [PATCH] Cleanup `alloc` / `std` imports for plonky2 (#1518)

* Cleanup alloc/std versions for plonky2

* Fix import for macro
---
 plonky2/examples/bench_recursion.rs           |  5 ++++
 plonky2/src/fri/mod.rs                        |  1 +
 plonky2/src/fri/oracle.rs                     |  4 ++--
 plonky2/src/fri/proof.rs                      |  4 ++--
 plonky2/src/fri/prover.rs                     |  1 +
 plonky2/src/fri/recursive_verifier.rs         |  4 ++--
 plonky2/src/fri/reduction_strategies.rs       |  4 ++--
 plonky2/src/fri/structure.rs                  |  1 +
 plonky2/src/fri/verifier.rs                   |  1 +
 plonky2/src/gadgets/arithmetic.rs             |  9 +++++---
 plonky2/src/gadgets/arithmetic_extension.rs   |  9 +++++---
 plonky2/src/gadgets/interpolation.rs          |  1 +
 plonky2/src/gadgets/lookup.rs                 |  4 ++--
 plonky2/src/gadgets/polynomial.rs             |  1 +
 plonky2/src/gadgets/random_access.rs          |  1 +
 plonky2/src/gadgets/range_check.rs            |  9 +++++---
 plonky2/src/gadgets/split_base.rs             |  5 ++--
 plonky2/src/gadgets/split_join.rs             |  9 +++++---
 plonky2/src/gates/arithmetic_base.rs          |  9 +++++---
 plonky2/src/gates/arithmetic_extension.rs     |  9 +++++---
 plonky2/src/gates/base_sum.rs                 |  5 ++--
 plonky2/src/gates/constant.rs                 |  5 ++--
 plonky2/src/gates/coset_interpolation.rs      | 10 +++++---
 plonky2/src/gates/exponentiation.rs           | 10 +++++---
 plonky2/src/gates/gate.rs                     |  8 +++----
 plonky2/src/gates/gate_testing.rs             |  4 ++--
 plonky2/src/gates/lookup.rs                   | 10 +++++---
 plonky2/src/gates/lookup_table.rs             | 14 +++++++----
 plonky2/src/gates/multiplication_extension.rs |  9 +++++---
 plonky2/src/gates/noop.rs                     |  4 ++--
 plonky2/src/gates/packed_util.rs              |  4 ++--
 plonky2/src/gates/poseidon.rs                 | 23 ++++++++-----------
 plonky2/src/gates/poseidon_mds.rs             | 10 +++++---
 plonky2/src/gates/public_input.rs             |  4 ++--
 plonky2/src/gates/random_access.rs            | 10 +++++---
 plonky2/src/gates/reducing.rs                 | 10 +++++---
 plonky2/src/gates/reducing_extension.rs       | 10 +++++---
 plonky2/src/gates/selectors.rs                |  4 ++--
 plonky2/src/hash/hash_types.rs                |  1 +
 plonky2/src/hash/hashing.rs                   |  2 +-
 plonky2/src/hash/keccak.rs                    |  4 ++--
 plonky2/src/hash/merkle_proofs.rs             |  4 ++--
 plonky2/src/hash/merkle_tree.rs               |  1 +
 plonky2/src/hash/path_compression.rs          |  4 ++--
 plonky2/src/hash/poseidon.rs                  | 10 +++-----
 plonky2/src/iop/challenger.rs                 |  4 ++--
 plonky2/src/iop/ext_target.rs                 |  1 +
 plonky2/src/iop/generator.rs                  | 11 +++++----
 plonky2/src/iop/target.rs                     |  1 +
 plonky2/src/iop/wire.rs                       |  1 +
 plonky2/src/iop/witness.rs                    |  4 ++--
 plonky2/src/lib.rs                            |  1 +
 plonky2/src/plonk/circuit_builder.rs          |  8 +++----
 plonky2/src/plonk/circuit_data.rs             |  7 +++---
 plonky2/src/plonk/config.rs                   |  4 ++--
 plonky2/src/plonk/copy_constraint.rs          |  1 +
 plonky2/src/plonk/get_challenges.rs           |  4 ++--
 plonky2/src/plonk/permutation_argument.rs     |  1 +
 plonky2/src/plonk/plonk_common.rs             |  4 ++--
 plonky2/src/plonk/proof.rs                    | 11 +++++----
 plonky2/src/plonk/prover.rs                   |  4 ++--
 plonky2/src/plonk/vanishing_poly.rs           |  4 ++--
 .../conditional_recursive_verifier.rs         |  1 +
 plonky2/src/recursion/cyclic_recursion.rs     |  1 +
 plonky2/src/recursion/dummy_circuit.rs        |  9 +++++---
 plonky2/src/util/context_tree.rs              |  9 +++++---
 plonky2/src/util/partial_products.rs          |  1 +
 plonky2/src/util/reducing.rs                  |  4 ++--
 .../util/serialization/gate_serialization.rs  |  5 +++-
 .../serialization/generator_serialization.rs  |  7 ++++--
 plonky2/src/util/serialization/mod.rs         |  8 +++----
 71 files changed, 235 insertions(+), 152 deletions(-)

diff --git a/plonky2/examples/bench_recursion.rs b/plonky2/examples/bench_recursion.rs
index 2e4c1ca3..8201c96d 100644
--- a/plonky2/examples/bench_recursion.rs
+++ b/plonky2/examples/bench_recursion.rs
@@ -3,11 +3,16 @@
 // put it in `src/bin/`, but then we wouldn't have access to
 // `[dev-dependencies]`.
 
+#[cfg(not(feature = "std"))]
 extern crate alloc;
+
+#[cfg(not(feature = "std"))]
 use alloc::sync::Arc;
 use core::num::ParseIntError;
 use core::ops::RangeInclusive;
 use core::str::FromStr;
+#[cfg(feature = "std")]
+use std::sync::Arc;
 
 use anyhow::{anyhow, Context as _, Result};
 use itertools::Itertools;
diff --git a/plonky2/src/fri/mod.rs b/plonky2/src/fri/mod.rs
index 207a2ea8..3445ada8 100644
--- a/plonky2/src/fri/mod.rs
+++ b/plonky2/src/fri/mod.rs
@@ -3,6 +3,7 @@
 //! It provides both a native implementation and an in-circuit version
 //! of the FRI verifier for recursive proof composition.
 
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 
 use serde::Serialize;
diff --git a/plonky2/src/fri/oracle.rs b/plonky2/src/fri/oracle.rs
index 8642a6c5..64dcbc60 100644
--- a/plonky2/src/fri/oracle.rs
+++ b/plonky2/src/fri/oracle.rs
@@ -1,5 +1,5 @@
-use alloc::format;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{format, vec::Vec};
 
 use itertools::Itertools;
 use plonky2_field::types::Field;
diff --git a/plonky2/src/fri/proof.rs b/plonky2/src/fri/proof.rs
index 71b62c71..edff1bea 100644
--- a/plonky2/src/fri/proof.rs
+++ b/plonky2/src/fri/proof.rs
@@ -1,5 +1,5 @@
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{vec, vec::Vec};
 
 use hashbrown::HashMap;
 use itertools::izip;
diff --git a/plonky2/src/fri/prover.rs b/plonky2/src/fri/prover.rs
index 378f1dae..4fb15614 100644
--- a/plonky2/src/fri/prover.rs
+++ b/plonky2/src/fri/prover.rs
@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 
 use plonky2_maybe_rayon::*;
diff --git a/plonky2/src/fri/recursive_verifier.rs b/plonky2/src/fri/recursive_verifier.rs
index da608242..47ae08f2 100644
--- a/plonky2/src/fri/recursive_verifier.rs
+++ b/plonky2/src/fri/recursive_verifier.rs
@@ -1,5 +1,5 @@
-use alloc::format;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{format, vec::Vec};
 
 use itertools::Itertools;
 
diff --git a/plonky2/src/fri/reduction_strategies.rs b/plonky2/src/fri/reduction_strategies.rs
index 6e575229..e7f5d799 100644
--- a/plonky2/src/fri/reduction_strategies.rs
+++ b/plonky2/src/fri/reduction_strategies.rs
@@ -1,5 +1,5 @@
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{vec, vec::Vec};
 
 use log::debug;
 use serde::Serialize;
diff --git a/plonky2/src/fri/structure.rs b/plonky2/src/fri/structure.rs
index 7d7436d5..81e462da 100644
--- a/plonky2/src/fri/structure.rs
+++ b/plonky2/src/fri/structure.rs
@@ -1,6 +1,7 @@
 //! Information about the structure of a FRI instance, in terms of the oracles and polynomials
 //! involved, and the points they are opened at.
 
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 use core::ops::Range;
 
diff --git a/plonky2/src/fri/verifier.rs b/plonky2/src/fri/verifier.rs
index f860ba30..89faa0f6 100644
--- a/plonky2/src/fri/verifier.rs
+++ b/plonky2/src/fri/verifier.rs
@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 
 use anyhow::{ensure, Result};
diff --git a/plonky2/src/gadgets/arithmetic.rs b/plonky2/src/gadgets/arithmetic.rs
index 9982628e..0e6806ff 100644
--- a/plonky2/src/gadgets/arithmetic.rs
+++ b/plonky2/src/gadgets/arithmetic.rs
@@ -1,6 +1,9 @@
-use alloc::string::{String, ToString};
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{
+    string::{String, ToString},
+    vec,
+    vec::Vec,
+};
 use core::borrow::Borrow;
 
 use crate::field::extension::Extendable;
diff --git a/plonky2/src/gadgets/arithmetic_extension.rs b/plonky2/src/gadgets/arithmetic_extension.rs
index 3c1deac3..649f4082 100644
--- a/plonky2/src/gadgets/arithmetic_extension.rs
+++ b/plonky2/src/gadgets/arithmetic_extension.rs
@@ -1,6 +1,9 @@
-use alloc::string::{String, ToString};
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{
+    string::{String, ToString},
+    vec,
+    vec::Vec,
+};
 use core::borrow::Borrow;
 
 use crate::field::extension::{Extendable, FieldExtension, OEF};
diff --git a/plonky2/src/gadgets/interpolation.rs b/plonky2/src/gadgets/interpolation.rs
index 6adbc427..39b048af 100644
--- a/plonky2/src/gadgets/interpolation.rs
+++ b/plonky2/src/gadgets/interpolation.rs
@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec;
 
 use plonky2_field::extension::Extendable;
diff --git a/plonky2/src/gadgets/lookup.rs b/plonky2/src/gadgets/lookup.rs
index 4ab765ba..0d9963a8 100644
--- a/plonky2/src/gadgets/lookup.rs
+++ b/plonky2/src/gadgets/lookup.rs
@@ -1,5 +1,5 @@
-use alloc::borrow::ToOwned;
-use alloc::vec;
+#[cfg(not(feature = "std"))]
+use alloc::{borrow::ToOwned, vec};
 
 use crate::field::extension::Extendable;
 use crate::gates::lookup::LookupGate;
diff --git a/plonky2/src/gadgets/polynomial.rs b/plonky2/src/gadgets/polynomial.rs
index d43d99c2..94fbe3b1 100644
--- a/plonky2/src/gadgets/polynomial.rs
+++ b/plonky2/src/gadgets/polynomial.rs
@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 
 use crate::field::extension::Extendable;
diff --git a/plonky2/src/gadgets/random_access.rs b/plonky2/src/gadgets/random_access.rs
index 85d2c714..0d99a3e9 100644
--- a/plonky2/src/gadgets/random_access.rs
+++ b/plonky2/src/gadgets/random_access.rs
@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 
 use crate::field::extension::Extendable;
diff --git a/plonky2/src/gadgets/range_check.rs b/plonky2/src/gadgets/range_check.rs
index 41af064a..9a66a6a6 100644
--- a/plonky2/src/gadgets/range_check.rs
+++ b/plonky2/src/gadgets/range_check.rs
@@ -1,6 +1,9 @@
-use alloc::string::{String, ToString};
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{
+    string::{String, ToString},
+    vec,
+    vec::Vec,
+};
 
 use crate::field::extension::Extendable;
 use crate::hash::hash_types::RichField;
diff --git a/plonky2/src/gadgets/split_base.rs b/plonky2/src/gadgets/split_base.rs
index a2c98ac7..1cdec862 100644
--- a/plonky2/src/gadgets/split_base.rs
+++ b/plonky2/src/gadgets/split_base.rs
@@ -1,6 +1,5 @@
-use alloc::string::String;
-use alloc::vec::Vec;
-use alloc::{format, vec};
+#[cfg(not(feature = "std"))]
+use alloc::{format, string::String, vec, vec::Vec};
 use core::borrow::Borrow;
 
 use itertools::Itertools;
diff --git a/plonky2/src/gadgets/split_join.rs b/plonky2/src/gadgets/split_join.rs
index 6901c8ca..2f35b94c 100644
--- a/plonky2/src/gadgets/split_join.rs
+++ b/plonky2/src/gadgets/split_join.rs
@@ -1,6 +1,9 @@
-use alloc::string::{String, ToString};
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{
+    string::{String, ToString},
+    vec,
+    vec::Vec,
+};
 
 use crate::field::extension::Extendable;
 use crate::gates::base_sum::BaseSumGate;
diff --git a/plonky2/src/gates/arithmetic_base.rs b/plonky2/src/gates/arithmetic_base.rs
index dfdd87e8..75489579 100644
--- a/plonky2/src/gates/arithmetic_base.rs
+++ b/plonky2/src/gates/arithmetic_base.rs
@@ -1,6 +1,9 @@
-use alloc::format;
-use alloc::string::{String, ToString};
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{
+    format,
+    string::{String, ToString},
+    vec::Vec,
+};
 
 use crate::field::extension::Extendable;
 use crate::field::packed::PackedField;
diff --git a/plonky2/src/gates/arithmetic_extension.rs b/plonky2/src/gates/arithmetic_extension.rs
index a19c6b4a..60eb912b 100644
--- a/plonky2/src/gates/arithmetic_extension.rs
+++ b/plonky2/src/gates/arithmetic_extension.rs
@@ -1,6 +1,9 @@
-use alloc::format;
-use alloc::string::{String, ToString};
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{
+    format,
+    string::{String, ToString},
+    vec::Vec,
+};
 use core::ops::Range;
 
 use crate::field::extension::{Extendable, FieldExtension};
diff --git a/plonky2/src/gates/base_sum.rs b/plonky2/src/gates/base_sum.rs
index 1d0f8f80..0f38415d 100644
--- a/plonky2/src/gates/base_sum.rs
+++ b/plonky2/src/gates/base_sum.rs
@@ -1,6 +1,5 @@
-use alloc::string::String;
-use alloc::vec::Vec;
-use alloc::{format, vec};
+#[cfg(not(feature = "std"))]
+use alloc::{format, string::String, vec, vec::Vec};
 use core::ops::Range;
 
 use crate::field::extension::Extendable;
diff --git a/plonky2/src/gates/constant.rs b/plonky2/src/gates/constant.rs
index 144e1ca3..cc62de7f 100644
--- a/plonky2/src/gates/constant.rs
+++ b/plonky2/src/gates/constant.rs
@@ -1,6 +1,5 @@
-use alloc::string::String;
-use alloc::vec::Vec;
-use alloc::{format, vec};
+#[cfg(not(feature = "std"))]
+use alloc::{format, string::String, vec, vec::Vec};
 
 use serde::{Deserialize, Serialize};
 
diff --git a/plonky2/src/gates/coset_interpolation.rs b/plonky2/src/gates/coset_interpolation.rs
index ab69f698..9911e927 100644
--- a/plonky2/src/gates/coset_interpolation.rs
+++ b/plonky2/src/gates/coset_interpolation.rs
@@ -1,6 +1,10 @@
-use alloc::string::{String, ToString};
-use alloc::vec::Vec;
-use alloc::{format, vec};
+#[cfg(not(feature = "std"))]
+use alloc::{
+    format,
+    string::{String, ToString},
+    vec,
+    vec::Vec,
+};
 use core::marker::PhantomData;
 use core::ops::Range;
 
diff --git a/plonky2/src/gates/exponentiation.rs b/plonky2/src/gates/exponentiation.rs
index 0011f011..2b7164e1 100644
--- a/plonky2/src/gates/exponentiation.rs
+++ b/plonky2/src/gates/exponentiation.rs
@@ -1,6 +1,10 @@
-use alloc::string::{String, ToString};
-use alloc::vec::Vec;
-use alloc::{format, vec};
+#[cfg(not(feature = "std"))]
+use alloc::{
+    format,
+    string::{String, ToString},
+    vec,
+    vec::Vec,
+};
 use core::marker::PhantomData;
 
 use crate::field::extension::Extendable;
diff --git a/plonky2/src/gates/gate.rs b/plonky2/src/gates/gate.rs
index cc8f7513..07de4fa3 100644
--- a/plonky2/src/gates/gate.rs
+++ b/plonky2/src/gates/gate.rs
@@ -1,11 +1,11 @@
-use alloc::string::String;
-use alloc::sync::Arc;
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{string::String, sync::Arc, vec, vec::Vec};
 use core::any::Any;
 use core::fmt::{Debug, Error, Formatter};
 use core::hash::{Hash, Hasher};
 use core::ops::Range;
+#[cfg(feature = "std")]
+use std::sync::Arc;
 
 use hashbrown::HashMap;
 use serde::{Serialize, Serializer};
diff --git a/plonky2/src/gates/gate_testing.rs b/plonky2/src/gates/gate_testing.rs
index 9a1f0f49..c71e96df 100644
--- a/plonky2/src/gates/gate_testing.rs
+++ b/plonky2/src/gates/gate_testing.rs
@@ -1,5 +1,5 @@
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{vec, vec::Vec};
 
 use anyhow::{ensure, Result};
 
diff --git a/plonky2/src/gates/lookup.rs b/plonky2/src/gates/lookup.rs
index 42b3bb92..23a0fd87 100644
--- a/plonky2/src/gates/lookup.rs
+++ b/plonky2/src/gates/lookup.rs
@@ -1,6 +1,10 @@
-use alloc::string::{String, ToString};
-use alloc::vec::Vec;
-use alloc::{format, vec};
+#[cfg(not(feature = "std"))]
+use alloc::{
+    format,
+    string::{String, ToString},
+    vec,
+    vec::Vec,
+};
 use core::usize;
 
 use itertools::Itertools;
diff --git a/plonky2/src/gates/lookup_table.rs b/plonky2/src/gates/lookup_table.rs
index ad01e092..9a4d08c8 100644
--- a/plonky2/src/gates/lookup_table.rs
+++ b/plonky2/src/gates/lookup_table.rs
@@ -1,8 +1,14 @@
-use alloc::string::{String, ToString};
-use alloc::sync::Arc;
-use alloc::vec::Vec;
-use alloc::{format, vec};
+#[cfg(not(feature = "std"))]
+use alloc::{
+    format,
+    string::{String, ToString},
+    sync::Arc,
+    vec,
+    vec::Vec,
+};
 use core::usize;
+#[cfg(feature = "std")]
+use std::sync::Arc;
 
 use itertools::Itertools;
 use keccak_hash::keccak;
diff --git a/plonky2/src/gates/multiplication_extension.rs b/plonky2/src/gates/multiplication_extension.rs
index 3f9fd8fe..143c854c 100644
--- a/plonky2/src/gates/multiplication_extension.rs
+++ b/plonky2/src/gates/multiplication_extension.rs
@@ -1,6 +1,9 @@
-use alloc::format;
-use alloc::string::{String, ToString};
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{
+    format,
+    string::{String, ToString},
+    vec::Vec,
+};
 use core::ops::Range;
 
 use crate::field::extension::{Extendable, FieldExtension};
diff --git a/plonky2/src/gates/noop.rs b/plonky2/src/gates/noop.rs
index 8752f380..54cb6422 100644
--- a/plonky2/src/gates/noop.rs
+++ b/plonky2/src/gates/noop.rs
@@ -1,5 +1,5 @@
-use alloc::string::String;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{string::String, vec::Vec};
 
 use crate::field::extension::Extendable;
 use crate::gates::gate::Gate;
diff --git a/plonky2/src/gates/packed_util.rs b/plonky2/src/gates/packed_util.rs
index 361eb3a2..32f1c37a 100644
--- a/plonky2/src/gates/packed_util.rs
+++ b/plonky2/src/gates/packed_util.rs
@@ -1,5 +1,5 @@
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{vec, vec::Vec};
 
 use crate::field::extension::Extendable;
 use crate::field::packable::Packable;
diff --git a/plonky2/src/gates/poseidon.rs b/plonky2/src/gates/poseidon.rs
index 3ba1b67b..be9a064e 100644
--- a/plonky2/src/gates/poseidon.rs
+++ b/plonky2/src/gates/poseidon.rs
@@ -1,6 +1,10 @@
-use alloc::string::{String, ToString};
-use alloc::vec::Vec;
-use alloc::{format, vec};
+#[cfg(not(feature = "std"))]
+use alloc::{
+    format,
+    string::{String, ToString},
+    vec,
+    vec::Vec,
+};
 use core::marker::PhantomData;
 
 use crate::field::extension::Extendable;
@@ -532,20 +536,13 @@ impl<F: RichField + Extendable<D> + Poseidon, const D: usize> SimpleGenerator<F,
 
 #[cfg(test)]
 mod tests {
-    #[cfg(not(feature = "std"))]
-    use alloc::{vec, vec::Vec};
-
     use anyhow::Result;
+    use plonky2_field::goldilocks_field::GoldilocksField;
 
-    use crate::field::goldilocks_field::GoldilocksField;
-    use crate::field::types::Field;
+    use super::*;
     use crate::gates::gate_testing::{test_eval_fns, test_low_degree};
-    use crate::gates::poseidon::PoseidonGate;
-    use crate::hash::poseidon::{Poseidon, SPONGE_WIDTH};
     use crate::iop::generator::generate_partial_witness;
-    use crate::iop::wire::Wire;
-    use crate::iop::witness::{PartialWitness, Witness, WitnessWrite};
-    use crate::plonk::circuit_builder::CircuitBuilder;
+    use crate::iop::witness::PartialWitness;
     use crate::plonk::circuit_data::CircuitConfig;
     use crate::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
 
diff --git a/plonky2/src/gates/poseidon_mds.rs b/plonky2/src/gates/poseidon_mds.rs
index 9b97cac7..9bd37e5e 100644
--- a/plonky2/src/gates/poseidon_mds.rs
+++ b/plonky2/src/gates/poseidon_mds.rs
@@ -1,6 +1,10 @@
-use alloc::string::{String, ToString};
-use alloc::vec::Vec;
-use alloc::{format, vec};
+#[cfg(not(feature = "std"))]
+use alloc::{
+    format,
+    string::{String, ToString},
+    vec,
+    vec::Vec,
+};
 use core::marker::PhantomData;
 use core::ops::Range;
 
diff --git a/plonky2/src/gates/public_input.rs b/plonky2/src/gates/public_input.rs
index 8be41d40..3a754c7c 100644
--- a/plonky2/src/gates/public_input.rs
+++ b/plonky2/src/gates/public_input.rs
@@ -1,5 +1,5 @@
-use alloc::string::String;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{string::String, vec::Vec};
 use core::ops::Range;
 
 use crate::field::extension::Extendable;
diff --git a/plonky2/src/gates/random_access.rs b/plonky2/src/gates/random_access.rs
index 59af01a5..086b119a 100644
--- a/plonky2/src/gates/random_access.rs
+++ b/plonky2/src/gates/random_access.rs
@@ -1,6 +1,10 @@
-use alloc::string::{String, ToString};
-use alloc::vec::Vec;
-use alloc::{format, vec};
+#[cfg(not(feature = "std"))]
+use alloc::{
+    format,
+    string::{String, ToString},
+    vec,
+    vec::Vec,
+};
 use core::marker::PhantomData;
 
 use itertools::Itertools;
diff --git a/plonky2/src/gates/reducing.rs b/plonky2/src/gates/reducing.rs
index c9daf538..5e0fc9f4 100644
--- a/plonky2/src/gates/reducing.rs
+++ b/plonky2/src/gates/reducing.rs
@@ -1,6 +1,10 @@
-use alloc::string::{String, ToString};
-use alloc::vec::Vec;
-use alloc::{format, vec};
+#[cfg(not(feature = "std"))]
+use alloc::{
+    format,
+    string::{String, ToString},
+    vec,
+    vec::Vec,
+};
 use core::ops::Range;
 
 use crate::field::extension::{Extendable, FieldExtension};
diff --git a/plonky2/src/gates/reducing_extension.rs b/plonky2/src/gates/reducing_extension.rs
index b1dc5e85..0ac3d9fd 100644
--- a/plonky2/src/gates/reducing_extension.rs
+++ b/plonky2/src/gates/reducing_extension.rs
@@ -1,6 +1,10 @@
-use alloc::string::{String, ToString};
-use alloc::vec::Vec;
-use alloc::{format, vec};
+#[cfg(not(feature = "std"))]
+use alloc::{
+    format,
+    string::{String, ToString},
+    vec,
+    vec::Vec,
+};
 use core::ops::Range;
 
 use crate::field::extension::{Extendable, FieldExtension};
diff --git a/plonky2/src/gates/selectors.rs b/plonky2/src/gates/selectors.rs
index be9a9da8..c75f2e22 100644
--- a/plonky2/src/gates/selectors.rs
+++ b/plonky2/src/gates/selectors.rs
@@ -1,5 +1,5 @@
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{vec, vec::Vec};
 use core::ops::Range;
 
 use serde::Serialize;
diff --git a/plonky2/src/hash/hash_types.rs b/plonky2/src/hash/hash_types.rs
index edd473c5..992f3619 100644
--- a/plonky2/src/hash/hash_types.rs
+++ b/plonky2/src/hash/hash_types.rs
@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 
 use anyhow::ensure;
diff --git a/plonky2/src/hash/hashing.rs b/plonky2/src/hash/hashing.rs
index f5fe1f1e..75b35a93 100644
--- a/plonky2/src/hash/hashing.rs
+++ b/plonky2/src/hash/hashing.rs
@@ -1,5 +1,5 @@
 //! Concrete instantiation of a hash function.
-
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 use core::fmt::Debug;
 
diff --git a/plonky2/src/hash/keccak.rs b/plonky2/src/hash/keccak.rs
index 281220f3..d3fa8c4b 100644
--- a/plonky2/src/hash/keccak.rs
+++ b/plonky2/src/hash/keccak.rs
@@ -1,5 +1,5 @@
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{vec, vec::Vec};
 use core::mem::size_of;
 
 use itertools::Itertools;
diff --git a/plonky2/src/hash/merkle_proofs.rs b/plonky2/src/hash/merkle_proofs.rs
index c848f66e..c8fd2492 100644
--- a/plonky2/src/hash/merkle_proofs.rs
+++ b/plonky2/src/hash/merkle_proofs.rs
@@ -1,5 +1,5 @@
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{vec, vec::Vec};
 
 use anyhow::{ensure, Result};
 use itertools::Itertools;
diff --git a/plonky2/src/hash/merkle_tree.rs b/plonky2/src/hash/merkle_tree.rs
index ab7e3958..10962727 100644
--- a/plonky2/src/hash/merkle_tree.rs
+++ b/plonky2/src/hash/merkle_tree.rs
@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 use core::mem::MaybeUninit;
 use core::slice;
diff --git a/plonky2/src/hash/path_compression.rs b/plonky2/src/hash/path_compression.rs
index bc2f23b0..517576bf 100644
--- a/plonky2/src/hash/path_compression.rs
+++ b/plonky2/src/hash/path_compression.rs
@@ -1,5 +1,5 @@
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{vec, vec::Vec};
 
 use hashbrown::HashMap;
 use num::Integer;
diff --git a/plonky2/src/hash/poseidon.rs b/plonky2/src/hash/poseidon.rs
index 2d357b40..204f04ba 100644
--- a/plonky2/src/hash/poseidon.rs
+++ b/plonky2/src/hash/poseidon.rs
@@ -1,8 +1,8 @@
 //! Implementation of the Poseidon hash function, as described in
 //! <https://eprint.iacr.org/2019/458.pdf>
 
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{vec, vec::Vec};
 use core::fmt::Debug;
 
 use unroll::unroll_for_loops;
@@ -753,11 +753,7 @@ impl<F: RichField> AlgebraicHasher<F> for PoseidonHash {
 
 #[cfg(test)]
 pub(crate) mod test_helpers {
-    #[cfg(not(feature = "std"))]
-    use alloc::vec::Vec;
-
-    use crate::field::types::Field;
-    use crate::hash::poseidon::{Poseidon, SPONGE_WIDTH};
+    use super::*;
 
     pub(crate) fn check_test_vectors<F: Field>(
         test_vectors: Vec<([u64; SPONGE_WIDTH], [u64; SPONGE_WIDTH])>,
diff --git a/plonky2/src/iop/challenger.rs b/plonky2/src/iop/challenger.rs
index d7b3c237..2daa7fdc 100644
--- a/plonky2/src/iop/challenger.rs
+++ b/plonky2/src/iop/challenger.rs
@@ -1,5 +1,5 @@
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{vec, vec::Vec};
 use core::marker::PhantomData;
 
 use crate::field::extension::{Extendable, FieldExtension};
diff --git a/plonky2/src/iop/ext_target.rs b/plonky2/src/iop/ext_target.rs
index c64d96e8..cc903557 100644
--- a/plonky2/src/iop/ext_target.rs
+++ b/plonky2/src/iop/ext_target.rs
@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 use core::ops::Range;
 
diff --git a/plonky2/src/iop/generator.rs b/plonky2/src/iop/generator.rs
index 1704b347..6cdd75dc 100644
--- a/plonky2/src/iop/generator.rs
+++ b/plonky2/src/iop/generator.rs
@@ -1,7 +1,10 @@
-use alloc::boxed::Box;
-use alloc::string::{String, ToString};
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{
+    boxed::Box,
+    string::{String, ToString},
+    vec,
+    vec::Vec,
+};
 use core::fmt::Debug;
 use core::marker::PhantomData;
 
diff --git a/plonky2/src/iop/target.rs b/plonky2/src/iop/target.rs
index 705941e0..f70d4c3d 100644
--- a/plonky2/src/iop/target.rs
+++ b/plonky2/src/iop/target.rs
@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 use core::ops::Range;
 
diff --git a/plonky2/src/iop/wire.rs b/plonky2/src/iop/wire.rs
index 435479ce..cfa69755 100644
--- a/plonky2/src/iop/wire.rs
+++ b/plonky2/src/iop/wire.rs
@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 use core::ops::Range;
 
diff --git a/plonky2/src/iop/witness.rs b/plonky2/src/iop/witness.rs
index cf74be51..40377aa3 100644
--- a/plonky2/src/iop/witness.rs
+++ b/plonky2/src/iop/witness.rs
@@ -1,5 +1,5 @@
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{vec, vec::Vec};
 
 use hashbrown::HashMap;
 use itertools::{zip_eq, Itertools};
diff --git a/plonky2/src/lib.rs b/plonky2/src/lib.rs
index 44bc2cf6..b0b6bfb4 100644
--- a/plonky2/src/lib.rs
+++ b/plonky2/src/lib.rs
@@ -2,6 +2,7 @@
 #![allow(clippy::needless_range_loop)]
 #![cfg_attr(not(feature = "std"), no_std)]
 
+#[cfg(not(feature = "std"))]
 pub extern crate alloc;
 
 /// Re-export of `plonky2_field`.
diff --git a/plonky2/src/plonk/circuit_builder.rs b/plonky2/src/plonk/circuit_builder.rs
index 4c2a5369..6df692fb 100644
--- a/plonky2/src/plonk/circuit_builder.rs
+++ b/plonky2/src/plonk/circuit_builder.rs
@@ -1,12 +1,10 @@
 //! Logic for building plonky2 circuits.
 
-use alloc::collections::BTreeMap;
-use alloc::sync::Arc;
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{collections::BTreeMap, sync::Arc, vec, vec::Vec};
 use core::cmp::max;
 #[cfg(feature = "std")]
-use std::time::Instant;
+use std::{collections::BTreeMap, sync::Arc, time::Instant};
 
 use hashbrown::{HashMap, HashSet};
 use itertools::Itertools;
diff --git a/plonky2/src/plonk/circuit_data.rs b/plonky2/src/plonk/circuit_data.rs
index d9847f4b..e4afb568 100644
--- a/plonky2/src/plonk/circuit_data.rs
+++ b/plonky2/src/plonk/circuit_data.rs
@@ -12,10 +12,11 @@
 //! The verifier data can similarly be extracted by calling [`CircuitData::verifier_data`].
 //! This is useful to allow even small devices to verify plonky2 proofs.
 
-use alloc::collections::BTreeMap;
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{collections::BTreeMap, vec, vec::Vec};
 use core::ops::{Range, RangeFrom};
+#[cfg(feature = "std")]
+use std::collections::BTreeMap;
 
 use anyhow::Result;
 use serde::Serialize;
diff --git a/plonky2/src/plonk/config.rs b/plonky2/src/plonk/config.rs
index 1ed40c40..ee5b69ff 100644
--- a/plonky2/src/plonk/config.rs
+++ b/plonky2/src/plonk/config.rs
@@ -6,8 +6,8 @@
 //! the Poseidon hash function both internally and natively, and one
 //! mixing Poseidon internally and truncated Keccak externally.
 
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{vec, vec::Vec};
 use core::fmt::Debug;
 
 use serde::de::DeserializeOwned;
diff --git a/plonky2/src/plonk/copy_constraint.rs b/plonky2/src/plonk/copy_constraint.rs
index ea92ec1c..cf7a6a19 100644
--- a/plonky2/src/plonk/copy_constraint.rs
+++ b/plonky2/src/plonk/copy_constraint.rs
@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::string::String;
 
 use crate::iop::target::Target;
diff --git a/plonky2/src/plonk/get_challenges.rs b/plonky2/src/plonk/get_challenges.rs
index ee6167b9..45d79f99 100644
--- a/plonky2/src/plonk/get_challenges.rs
+++ b/plonky2/src/plonk/get_challenges.rs
@@ -1,5 +1,5 @@
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{vec, vec::Vec};
 
 use hashbrown::HashSet;
 
diff --git a/plonky2/src/plonk/permutation_argument.rs b/plonky2/src/plonk/permutation_argument.rs
index a0dd5770..312f3e99 100644
--- a/plonky2/src/plonk/permutation_argument.rs
+++ b/plonky2/src/plonk/permutation_argument.rs
@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 
 use hashbrown::HashMap;
diff --git a/plonky2/src/plonk/plonk_common.rs b/plonky2/src/plonk/plonk_common.rs
index ca8ea919..170bfa17 100644
--- a/plonky2/src/plonk/plonk_common.rs
+++ b/plonky2/src/plonk/plonk_common.rs
@@ -1,7 +1,7 @@
 //! Utility methods and constants for Plonk.
 
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{vec, vec::Vec};
 
 use crate::field::extension::Extendable;
 use crate::field::packed::PackedField;
diff --git a/plonky2/src/plonk/proof.rs b/plonky2/src/plonk/proof.rs
index de82746a..a9151a9f 100644
--- a/plonky2/src/plonk/proof.rs
+++ b/plonky2/src/plonk/proof.rs
@@ -4,8 +4,8 @@
 //! [`CompressedProof`] or [`CompressedProofWithPublicInputs`] formats.
 //! The latter can be directly passed to a verifier to assert its correctness.
 
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{vec, vec::Vec};
 
 use anyhow::ensure;
 use plonky2_maybe_rayon::*;
@@ -452,21 +452,22 @@ impl<const D: usize> OpeningSetTarget<D> {
 #[cfg(test)]
 mod tests {
     #[cfg(not(feature = "std"))]
-    use alloc::{sync::Arc, vec};
+    use alloc::sync::Arc;
     #[cfg(feature = "std")]
     use std::sync::Arc;
 
     use anyhow::Result;
     use itertools::Itertools;
+    use plonky2_field::types::Sample;
 
-    use crate::field::types::Sample;
+    use super::*;
     use crate::fri::reduction_strategies::FriReductionStrategy;
     use crate::gates::lookup_table::LookupTable;
     use crate::gates::noop::NoopGate;
     use crate::iop::witness::PartialWitness;
     use crate::plonk::circuit_builder::CircuitBuilder;
     use crate::plonk::circuit_data::CircuitConfig;
-    use crate::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
+    use crate::plonk::config::PoseidonGoldilocksConfig;
     use crate::plonk::verifier::verify;
 
     #[test]
diff --git a/plonky2/src/plonk/prover.rs b/plonky2/src/plonk/prover.rs
index 153610dc..400845bc 100644
--- a/plonky2/src/plonk/prover.rs
+++ b/plonky2/src/plonk/prover.rs
@@ -1,7 +1,7 @@
 //! plonky2 prover implementation.
 
-use alloc::vec::Vec;
-use alloc::{format, vec};
+#[cfg(not(feature = "std"))]
+use alloc::{format, vec, vec::Vec};
 use core::cmp::min;
 use core::mem::swap;
 
diff --git a/plonky2/src/plonk/vanishing_poly.rs b/plonky2/src/plonk/vanishing_poly.rs
index e3ddcf5b..9eb7252e 100644
--- a/plonky2/src/plonk/vanishing_poly.rs
+++ b/plonky2/src/plonk/vanishing_poly.rs
@@ -1,5 +1,5 @@
-use alloc::vec::Vec;
-use alloc::{format, vec};
+#[cfg(not(feature = "std"))]
+use alloc::{format, vec, vec::Vec};
 use core::cmp::min;
 
 use plonky2_field::polynomial::PolynomialCoeffs;
diff --git a/plonky2/src/recursion/conditional_recursive_verifier.rs b/plonky2/src/recursion/conditional_recursive_verifier.rs
index a35b46ea..43bef589 100644
--- a/plonky2/src/recursion/conditional_recursive_verifier.rs
+++ b/plonky2/src/recursion/conditional_recursive_verifier.rs
@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 
 use itertools::Itertools;
diff --git a/plonky2/src/recursion/cyclic_recursion.rs b/plonky2/src/recursion/cyclic_recursion.rs
index 172c0826..7be55417 100644
--- a/plonky2/src/recursion/cyclic_recursion.rs
+++ b/plonky2/src/recursion/cyclic_recursion.rs
@@ -1,5 +1,6 @@
 #![allow(clippy::int_plus_one)] // Makes more sense for some inequalities below.
 
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 
 use anyhow::{ensure, Result};
diff --git a/plonky2/src/recursion/dummy_circuit.rs b/plonky2/src/recursion/dummy_circuit.rs
index ee73105a..501a475f 100644
--- a/plonky2/src/recursion/dummy_circuit.rs
+++ b/plonky2/src/recursion/dummy_circuit.rs
@@ -1,6 +1,9 @@
-use alloc::string::{String, ToString};
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{
+    string::{String, ToString},
+    vec,
+    vec::Vec,
+};
 
 use hashbrown::HashMap;
 use plonky2_field::extension::Extendable;
diff --git a/plonky2/src/util/context_tree.rs b/plonky2/src/util/context_tree.rs
index a0a69971..2e70fd61 100644
--- a/plonky2/src/util/context_tree.rs
+++ b/plonky2/src/util/context_tree.rs
@@ -1,6 +1,9 @@
-use alloc::string::{String, ToString};
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{
+    string::{String, ToString},
+    vec,
+    vec::Vec,
+};
 
 use log::{log, Level};
 
diff --git a/plonky2/src/util/partial_products.rs b/plonky2/src/util/partial_products.rs
index e195af1a..1ceea7ca 100644
--- a/plonky2/src/util/partial_products.rs
+++ b/plonky2/src/util/partial_products.rs
@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 use core::iter;
 
diff --git a/plonky2/src/util/reducing.rs b/plonky2/src/util/reducing.rs
index e1ba397b..b99da32e 100644
--- a/plonky2/src/util/reducing.rs
+++ b/plonky2/src/util/reducing.rs
@@ -1,5 +1,5 @@
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{vec, vec::Vec};
 use core::borrow::Borrow;
 
 use crate::field::extension::{Extendable, FieldExtension};
diff --git a/plonky2/src/util/serialization/gate_serialization.rs b/plonky2/src/util/serialization/gate_serialization.rs
index c5763fb0..8b10da07 100644
--- a/plonky2/src/util/serialization/gate_serialization.rs
+++ b/plonky2/src/util/serialization/gate_serialization.rs
@@ -1,6 +1,9 @@
 //! A module to help with GateRef serialization
 
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
+#[cfg(feature = "std")]
+use std::vec::Vec; // For macros below
 
 use plonky2_field::extension::Extendable;
 
@@ -76,7 +79,7 @@ macro_rules! impl_gate_serializer {
 
         fn write_gate(
             &self,
-            buf: &mut $crate::alloc::vec::Vec<u8>,
+            buf: &mut $crate::util::serialization::gate_serialization::Vec<u8>,
             gate: &$crate::gates::gate::GateRef<F, D>,
             common: &$crate::plonk::circuit_data::CommonCircuitData<F, D>,
         ) -> $crate::util::serialization::IoResult<()> {
diff --git a/plonky2/src/util/serialization/generator_serialization.rs b/plonky2/src/util/serialization/generator_serialization.rs
index bad24ceb..6ede0020 100644
--- a/plonky2/src/util/serialization/generator_serialization.rs
+++ b/plonky2/src/util/serialization/generator_serialization.rs
@@ -1,6 +1,9 @@
 //! A module to help with WitnessGeneratorRef serialization
 
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+pub use alloc::vec::Vec;
+#[cfg(feature = "std")]
+pub use std::vec::Vec; // For macros below
 
 use plonky2_field::extension::Extendable;
 
@@ -80,7 +83,7 @@ macro_rules! impl_generator_serializer {
 
         fn write_generator(
             &self,
-            buf: &mut $crate::alloc::vec::Vec<u8>,
+            buf: &mut $crate::util::serialization::generator_serialization::Vec<u8>,
             generator: &$crate::iop::generator::WitnessGeneratorRef<F, D>,
             common: &$crate::plonk::circuit_data::CommonCircuitData<F, D>,
         ) -> $crate::util::serialization::IoResult<()> {
diff --git a/plonky2/src/util/serialization/mod.rs b/plonky2/src/util/serialization/mod.rs
index 94551bdf..393db6c6 100644
--- a/plonky2/src/util/serialization/mod.rs
+++ b/plonky2/src/util/serialization/mod.rs
@@ -4,14 +4,14 @@ pub mod generator_serialization;
 #[macro_use]
 pub mod gate_serialization;
 
-use alloc::collections::BTreeMap;
-use alloc::sync::Arc;
-use alloc::vec;
-use alloc::vec::Vec;
+#[cfg(not(feature = "std"))]
+use alloc::{collections::BTreeMap, sync::Arc, vec, vec::Vec};
 use core::convert::Infallible;
 use core::fmt::{Debug, Display, Formatter};
 use core::mem::size_of;
 use core::ops::Range;
+#[cfg(feature = "std")]
+use std::{collections::BTreeMap, sync::Arc};
 
 pub use gate_serialization::default::DefaultGateSerializer;
 pub use gate_serialization::GateSerializer;