Cleanup alloc / std imports for plonky2 (#1518)

* Cleanup alloc/std versions for plonky2

* Fix import for macro

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;

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;

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;

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;

plonky2/src/fri/prover.rs

@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 
 use plonky2_maybe_rayon::*;

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;
 

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;

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;
 

plonky2/src/fri/verifier.rs

@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 
 use anyhow::{ensure, Result};

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;

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};

plonky2/src/gadgets/interpolation.rs

@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec;
 
 use plonky2_field::extension::Extendable;

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;

plonky2/src/gadgets/polynomial.rs

@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 
 use crate::field::extension::Extendable;

plonky2/src/gadgets/random_access.rs

@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 
 use crate::field::extension::Extendable;

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;

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;

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;

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;

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};

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;

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};
 

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;
 

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;

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};

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};
 

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;

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;

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};

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;

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;

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};
 

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;
 

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;

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;

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};

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};

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;

plonky2/src/hash/hash_types.rs

@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 
 use anyhow::ensure;

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;
 

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;

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;

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;

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;

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])>,

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};

plonky2/src/iop/ext_target.rs

@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 use core::ops::Range;
 

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;
 

plonky2/src/iop/target.rs

@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 use core::ops::Range;
 

plonky2/src/iop/wire.rs

@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 use core::ops::Range;
 

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};

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`.

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;

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;

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;

plonky2/src/plonk/copy_constraint.rs

@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::string::String;
 
 use crate::iop::target::Target;

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;
 

plonky2/src/plonk/permutation_argument.rs

@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 
 use hashbrown::HashMap;

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;

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]

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;
 

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;

plonky2/src/recursion/conditional_recursive_verifier.rs

@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 
 use itertools::Itertools;

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};

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;

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};
 

plonky2/src/util/partial_products.rs

@@ -1,3 +1,4 @@
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 use core::iter;
 

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};

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<()> {

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<()> {

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;