content: package dk.alexandra.fresco.lib.mimc;
  
  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.DefaultCollections;
  import dk.alexandra.fresco.lib.common.collections.Matrix;
  import java.math.BigInteger;
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Map.Entry;
  
  public class MiMCAggregation implements Computation<Matrix<DRes<SInt>>, ProtocolBuilderNumeric> {
  
    private final DRes<Matrix<DRes<SInt>>> values;
    private final int groupColIdx;
    private final int aggColIdx;
  
    /**
     * Performs a SQL-like group-by sum operation. Groups rows by column <code>groupColIdx</code> and
     * sums values in resulting groups in column <code>aggColIdx</code>. <br>
     * NOTE: this particular implementation leaks equality of values in column <code>groupColIdx
     * </code> and the size of the result.
     *
     * @param values rows to be aggregated
     * @param groupColIdx column to group by
     * @param aggColIdx column to aggregate
     */
    public MiMCAggregation(DRes<Matrix<DRes<SInt>>> values, int groupColIdx, int aggColIdx) {
      super();
      this.values = values;
      this.groupColIdx = groupColIdx;
      this.aggColIdx = aggColIdx;
    }
  
    private Matrix<DRes<SInt>> toMatrix(
        Map<BigInteger, DRes<SInt>> groupedByCipher, Map<BigInteger, DRes<SInt>> cipherToShare) {
      ArrayList<ArrayList<DRes<SInt>>> result = new ArrayList<>(groupedByCipher.size());
      for (Entry<BigInteger, DRes<SInt>> keyAndValues : groupedByCipher.entrySet()) {
        ArrayList<DRes<SInt>> row = new ArrayList<>(2);
        row.add(cipherToShare.get(keyAndValues.getKey()));
        row.add(keyAndValues.getValue());
        result.add(row);
      }
      return new Matrix<>(result.size(), 2, result);
    }
  
    @Override
    public DRes<Matrix<DRes<SInt>>> buildComputation(ProtocolBuilderNumeric builder) {
      // shuffle input
      DRes<Matrix<DRes<SInt>>> shuffled = Collections.using(builder).shuffle(values);
      // generate encryption key
      DRes<SInt> mimcKey = builder.numeric().randomElement();
      return builder
          .par(
              par -> {
                Matrix<DRes<SInt>> inputRows = shuffled.out();
                // encrypt values in group by column and reveal cipher text
                List<TripleWithCipher> ciphers = new ArrayList<>(inputRows.getHeight());
                for (final ArrayList<DRes<SInt>> row : inputRows.getRows()) {
                  DRes<SInt> groupBy = row.get(groupColIdx);
                  DRes<SInt> aggOn = row.get(aggColIdx);
                  // encrypt and open groupBy
                  DRes<BigInteger> openedCipher =
                      par.seq(
                          (seq -> {
                            // TODO: encryption should be on a directory
                            DRes<SInt> cipherText = seq.seq(new MiMCEncryption(groupBy, mimcKey));
                            return seq.numeric().open(cipherText);
                          }));
                  ciphers.add(new TripleWithCipher(groupBy, aggOn, openedCipher));
                }
                return () -> ciphers;
              })
          .seq(
              (seq, triples) -> {
                // use cipher texts to perform aggregation "in-the-clear"
                Map<BigInteger, DRes<SInt>> groupedByCipher = new HashMap<>();
                Map<BigInteger, DRes<SInt>> cipherToShare = new HashMap<>();
  
                for (TripleWithCipher triple : triples) {
                  BigInteger cipher = triple.getCipher().out();
                  DRes<SInt> key = triple.getKey();
                  DRes<SInt> value = triple.getValue();
  
                  if (!groupedByCipher.containsKey(cipher)) {
                    groupedByCipher.put(cipher, value);
                    cipherToShare.put(cipher, key);
                  } else {
                    DRes<SInt> subTotal = seq.numeric().add(groupedByCipher.get(cipher), value);
                    groupedByCipher.put(cipher, subTotal);
                  }
                }
                return () -> toMatrix(groupedByCipher, cipherToShare);
              });
    }
  
    private class TripleWithCipher extends Triple<DRes<SInt>, DRes<SInt>, DRes<BigInteger>> {
  
      public TripleWithCipher(DRes<SInt> key, DRes<SInt> value, DRes<BigInteger> cipher) {
        super(key, value, cipher);
      }
    }
  
    private class Triple<K, V, C> {
  
      private final K key;
      private final V value;
      private final C cipher;
  
      public Triple(K key, V value, C cipher) {
        this.key = key;
        this.value = value;
        this.cipher = cipher;
      }
  
      public K getKey() {
        return key;
      }
  
      public V getValue() {
        return value;
      }
  
      public C getCipher() {
        return cipher;
      }
    }
  }
  