Prepare :

Create a HashSet of the Public Inputs Indices. (for faster lookup later)

 // Extract ACIR public input indices set from the noir program
 let acir_public_inputs_indices_set: HashSet<u32> =
 main.public_inputs().indices().iter().cloned().collect();

Modify the layer prep for w1 and w2 witness split to make sure all public inputs (which are of format WitnessBuilder::ACIR ) are in w1

 pub fn split_and_prepare_layers(
        witness_builders: &[WitnessBuilder],
        r1cs: R1CS,
        witness_map: Vec<Option<NonZeroU32>>,
        acir_public_inputs_indices_set: HashSet<u32>,
    ) -> (SplitWitnessBuilders, R1CS, Vec<Option<NonZeroU32>>, usize) {

pub fn split_builders(&self, acir_public_inputs_indices_set: HashSet<u32>) -> (Vec<usize>, Vec<usize>) {

Trivial Case : If there are no challenge builders, just rearrange w1 to put public input indices at front.

if challenge_builders.is_empty() {
    let w1_indices = self.rearrange_w1(
        (0..builder_count).collect(),
        &acir_public_inputs_indices_set,
    );
    return (w1_indices, Vec::new());
}

In a non trivial case : when challenge builders are present.

ACIR by default are put into w1 because they don't depend on any challenges. So, post the formation of mandatory_w1 (w1_set in the below code) and mandatory_w2 , make sure that in the rearranging code, all public indices go into w1

// If free builder writes a public witness, add it to w1_set.
if let WitnessBuilder::Acir(_, acir_idx) = &self.witness_builders[idx] {
    if acir_public_inputs_indices_set.contains(&(*acir_idx as u32)) {
        w1_set.insert(idx);
        w1_witness_count += witness_count;
        continue;
    }
}
....
w1_indices = self.rearrange_w1(w1_indices, &acir_public_inputs_indices_set);

Final result : witness builder w1 contains [0, builder which are associated with public inputs, then rest]

Prove :