content: package dk.alexandra.fresco.lib.common.collections.permute;
  
  import dk.alexandra.fresco.framework.DRes;
  import dk.alexandra.fresco.framework.builder.Computation;
  import dk.alexandra.fresco.framework.builder.numeric.ProtocolBuilderNumeric;
  import dk.alexandra.fresco.framework.value.SInt;
  import dk.alexandra.fresco.lib.common.collections.Collections;
  import dk.alexandra.fresco.lib.common.collections.Matrix;
  import java.util.ArrayList;
  
  public class PermuteRows implements Computation<Matrix<DRes<SInt>>, ProtocolBuilderNumeric> {
  
    private final DRes<Matrix<DRes<SInt>>> values;
    private final int[] idxPerm;
    private final int permProviderPid;
    private final boolean isPermProvider;
    private final WaksmanUtils wutils;
    // not final as this will be set during protocol execution
    private Matrix<DRes<SInt>> cbits;
  
    /**
     * Constructs a new PermuteRows computation.
     * 
     * @param values rows to permute
     * @param idxPerm encodes the desired permutation by supplying for each index a new index
     * @param permProviderPid the ID of the party choosing permutation
     * @param isPermProvider flag that indicates if I provide permutation
     */
    public PermuteRows(DRes<Matrix<DRes<SInt>>> values, int[] idxPerm, int permProviderPid,
        boolean isPermProvider) {
      super();
      this.wutils = new WaksmanUtils();
      this.values = values;
      this.idxPerm = idxPerm;
      this.isPermProvider = isPermProvider;
      this.permProviderPid = permProviderPid;
    }
  
    private Matrix<DRes<SInt>> reroute(Matrix<DRes<SInt>> roundInputs, int numRows, int numCols,
        int numSwapperRows, int numSwapperCols, int colIdx) {
      // this will store the re-arranged result
      ArrayList<ArrayList<DRes<SInt>>> rearranged = new ArrayList<>(numRows);
      // pre-initialize array list since we will be setting elements at indeces
      // as opposed to adding
      for (int x = 0; x < numRows; x++) {
        rearranged.add(null);
      }
      // determines if we are at a column before the center of the
      // network or after (fanning in vs fanning out)
      boolean inward = (colIdx < (int) numSwapperCols / 2);
      if (!inward) {
        colIdx = (numSwapperCols / 2) - colIdx % (numSwapperCols / 2) - 1;
      }
      int numPerms = (int) Math.pow(2, colIdx + 1);
      int elsPerPerm = (int) numRows / numPerms;
      for (int permPairIdx = 0; permPairIdx < numPerms - 1; permPairIdx += 2) {
        int topPermStart = permPairIdx * elsPerPerm;
        int bottomPermStart = (permPairIdx + 1) * elsPerPerm;
        int nextFreeTop = topPermStart;
        int nextFreeBottom = bottomPermStart;
        // re-arrange elements
        if (inward) {
          // we're haven't reached the center of the network yet
          // that is, we are fanning in
          for (int i = 0; i < elsPerPerm * 2 - 1; i += 2) {
            int inputIdx = topPermStart + i;
            rearranged.set(nextFreeTop++, roundInputs.getRow(inputIdx));
            rearranged.set(nextFreeBottom++, roundInputs.getRow(inputIdx + 1));
          }
        } else {
          // fanning out
          for (int i = 0; i < elsPerPerm * 2 - 1; i += 2) {
            int inputIdx = topPermStart + i;
            rearranged.set(inputIdx, roundInputs.getRow(nextFreeTop++));
            rearranged.set(inputIdx + 1, roundInputs.getRow(nextFreeBottom++));
          }
        }
      }
      return new Matrix<>(numRows, numCols, rearranged);
    }
  
    @Override
    public DRes<Matrix<DRes<SInt>>> buildComputation(ProtocolBuilderNumeric builder)
        throws UnsupportedOperationException {
      // sort of hacky: need to unwrap but will later used wrapped
      Matrix<DRes<SInt>> valuesOut = values.out();
      // determine dimensions of waksman network
      final int numSwapperRows = wutils.getNumRowsRequired(valuesOut.getHeight());
      final int numSwapperCols = wutils.getNumColsRequired(valuesOut.getHeight());
      // dimensions of row matrix
      final int numRows = valuesOut.getHeight();
      final int numCols = valuesOut.getWidth();
      // number of rounds will be number of swapper columns
      final int numRounds = numSwapperCols;
  
      if (numRounds == 0) {
        // in case of empty input, just return
        return values;
      }
  
      // throw if size is not power of 2
      if (!wutils.isPow2(numRows)) {
        throw new UnsupportedOperationException("input size must be power of 2");
      }
  
      // non-empty input, i.e., main protocol
      return builder.seq(seq -> {
        Collections collections = Collections.using(seq);
        if (isPermProvider) {
          return collections.closeMatrix(wutils.setControlBits(idxPerm), permProviderPid);
        } else {
          return collections.closeMatrix(numSwapperRows, numSwapperCols, permProviderPid);
        }
      }).seq((seq, bits) -> {
        // set control bits
        cbits = bits;
        // initiate loop
        return new IterationState(0, values);
      }).whileLoop((state) -> state.round < numRounds - 1, (seq, state) -> {
        Collections collections = Collections.using(seq);
  
        // apply swapper gates for this round
        DRes<Matrix<DRes<SInt>>> swapped =
            collections.swapNeighborsIf(() -> cbits.getColumn(state.round), state.intermediate);
        // re-arrange values for next round (based solely on waksman network topology)
        // this is NOT input-dependent!
        return new IterationState(state.round + 1, () -> reroute(swapped.out(), numRows, numCols,
            numSwapperRows, numSwapperCols, state.round));
      }).seq((seq, state) -> {
        Collections collections = Collections.using(seq);
        // Apply last column of swapper gates
        return collections.swapNeighborsIf(() -> cbits.getColumn(state.round),
            state.intermediate);
      });
    }
  
    private static final class IterationState implements DRes<IterationState> {
  
      private final int round;
      private final DRes<Matrix<DRes<SInt>>> intermediate;
  
      private IterationState(int round, DRes<Matrix<DRes<SInt>>> intermediate) {
        this.round = round;
        this.intermediate = intermediate;
      }
  
      @Override
      public IterationState out() {
        return this;
      }
    }
  
  }
  