content: package dk.alexandra.fresco.lib.common.collections.sort;
  
  import dk.alexandra.fresco.framework.DRes;
  import dk.alexandra.fresco.framework.builder.ComputationParallel;
  import dk.alexandra.fresco.framework.builder.ProtocolBuilderImpl;
  import dk.alexandra.fresco.framework.builder.binary.ProtocolBuilderBinary;
  import dk.alexandra.fresco.framework.builder.numeric.ProtocolBuilderNumeric;
  import dk.alexandra.fresco.framework.util.Pair;
  import dk.alexandra.fresco.framework.value.SBool;
  import dk.alexandra.fresco.framework.value.SInt;
  import dk.alexandra.fresco.lib.common.compare.BinaryComparison;
  import dk.alexandra.fresco.lib.common.compare.Comparison;
  import dk.alexandra.fresco.lib.common.math.AdvancedBinary;
  import dk.alexandra.fresco.lib.common.math.AdvancedNumeric;
  import java.util.ArrayList;
  import java.util.List;
  
  /**
   * @param <KeyT> The type of keys
   * @param <ValueT> The type of elements in the payload which will be a list of elements.
   * @param <ConditionT> The type of element representing a boolean value, so it should have a canonical
   *                    interpretation as a boolean 0/1 value. Used as condition in conditional swap.
   * @param <BuilderT> The type of {@link ProtocolBuilderImpl} used.
   */
  public class KeyedCompareAndSwap<KeyT, ValueT,
      ConditionT, BuilderT extends ProtocolBuilderImpl<BuilderT>> implements
      ComputationParallel<List<Pair<KeyT, List<ValueT>>>, BuilderT> {
  
    private final Compare<KeyT, ConditionT, BuilderT> comparison;
    private final ConditionalSelect<ConditionT, KeyT, BuilderT> condSelectKey;
    private final ConditionalSelect<ConditionT, List<ValueT>, BuilderT> condSelectValue;
    private final Addition<KeyT, BuilderT> additionKey;
    private final Addition<List<ValueT>, BuilderT> additionValue;
    private final Addition<KeyT, BuilderT> subtractionKey;
    private final Addition<List<ValueT>, BuilderT> subtractionValue;
    private final KeyT leftKey;
    private final KeyT rightKey;
    private final List<ValueT> leftValue;
    private final List<ValueT> rightValue;
    private KeyT xorKey;
    private List<ValueT> xorValue;
  
    /**
     *
     * @param leftKeyAndValue The left key value pair
     * @param rightKeyAndValue The right key value pair
     * @param comparison A function comparing two elements of type KeyT. Should return TRUE if the first value is greater than the second and return an element of type ConditionT.
     * @param additionKey A function adding two elements of type KeyT.
     * @param subtractionKey A function subtracting two elements of type KeyT.
     * @param additionValue A function adding two elements of type ValueT.
     * @param subtractionValue A function subtracting two elements of type ValueT.
     * @param condSelectKey A conditional select function for elements of type KeyT. Returns the first
     *                      value if the condition is TRUE and the second otherwise.
     * @param condSelectValue
     */
    private KeyedCompareAndSwap(
        Pair<KeyT, List<ValueT>> leftKeyAndValue,
        Pair<KeyT, List<ValueT>> rightKeyAndValue,
        Compare<KeyT, ConditionT, BuilderT> comparison,
        Addition<KeyT, BuilderT> additionKey,
        Addition<KeyT, BuilderT> subtractionKey,
        Addition<ValueT, BuilderT> additionValue,
        Addition<ValueT, BuilderT> subtractionValue,
        ConditionalSelect<ConditionT, KeyT, BuilderT> condSelectKey,
        ConditionalSelect<ConditionT, ValueT, BuilderT> condSelectValue) {
      this.leftKey = leftKeyAndValue.getFirst();
      this.leftValue = leftKeyAndValue.getSecond();
      this.rightKey = rightKeyAndValue.getFirst();
      this.rightValue = rightKeyAndValue.getSecond();
  
      this.comparison = comparison;
      this.additionKey = additionKey;
      this.subtractionKey = subtractionKey;
      this.additionValue = Addition.forLists(additionValue);
      this.subtractionValue = Addition.forLists(subtractionValue);
      this.condSelectKey = condSelectKey;
      this.condSelectValue = ConditionalSelect.forLists(condSelectValue);
    }
  
    /**
     * Constructs a protocol producer for the keyed compare and swap protocol. This protocol will
     * compare the keys of two key-value pairs and produce a list of pairs so that the first pair has
     * the largest key.
     *
     * @param leftKeyAndValue  the key and value of the left pair
     * @param rightKeyAndValue the key and value of the right pair
     */
    public static KeyedCompareAndSwap<List<DRes<SBool>>,
        DRes<SBool>, DRes<SBool>, ProtocolBuilderBinary> binary(
        Pair<List<DRes<SBool>>, List<DRes<SBool>>> leftKeyAndValue,
        Pair<List<DRes<SBool>>, List<DRes<SBool>>> rightKeyAndValue) {
  
      return new KeyedCompareAndSwap<List<DRes<SBool>>,
          DRes<SBool>, DRes<SBool>, ProtocolBuilderBinary>(leftKeyAndValue, rightKeyAndValue,
          (a, b, builder) -> BinaryComparison
              .using(builder).greaterThan(a, b),
          Addition
              .forLists((ai, bi, builder) -> builder.binary().xor(ai, bi)),
          Addition
              .forLists((ai, bi, builder) -> builder.binary().xor(ai, bi)),
          (ai, bi, builder) -> builder.binary().xor(ai, bi),
          (ai, bi, builder) -> builder.binary().xor(ai, bi),
          ConditionalSelect.forLists(
              (c, ai, bi, builder) -> AdvancedBinary.using(builder)
                  .condSelect(c, ai, bi)),
          (c, ai, bi, builder) -> AdvancedBinary.using(builder)
              .condSelect(c, ai, bi));
    }
  
    /**
     * Constructs a protocol producer for the keyed compare and swap protocol. This protocol will
     * compare the keys of two key-value pairs and produce a list of pairs so that the first pair has
     * the largest key.
     *
     * @param leftKeyAndValue  the key and value of the left pair
     * @param rightKeyAndValue the key and value of the right pair
     */
    public static KeyedCompareAndSwap<DRes<SInt>,
        DRes<SInt>, DRes<SInt>, ProtocolBuilderNumeric> numeric(
        Pair<DRes<SInt>, List<DRes<SInt>>> leftKeyAndValue,
        Pair<DRes<SInt>, List<DRes<SInt>>> rightKeyAndValue) {
  
      return new KeyedCompareAndSwap<>(leftKeyAndValue, rightKeyAndValue,
          (a, b, builder) -> Comparison
              .using(builder).compareLT(b, a),
          (a, b, builder) -> builder.numeric().add(a, b),
          (a, b, builder) -> builder.numeric().sub(a, b),
          (a, b, builder) -> builder.numeric().add(a, b),
          (a, b, builder) -> builder.numeric().sub(a, b),
          (c, a, b, builder) -> AdvancedNumeric.using(builder)
              .condSelect(c, a, b),
          (c, a, b, builder) -> AdvancedNumeric.using(builder)
              .condSelect(c, a, b));
    }
  
    @Override
    public DRes<List<Pair<KeyT, List<ValueT>>>> buildComputation(
        BuilderT builder) {
      return builder.par(seq -> {
        ConditionT c = comparison.compare(leftKey, rightKey, seq);
        xorKey = additionKey.add(leftKey, rightKey, seq);
        xorValue = additionValue.add(leftValue, rightValue, seq);
        return () -> c;
      }).par((par, c) -> {
        List<ValueT> firstValue = condSelectValue.apply(
            c, leftValue, rightValue, par);
        KeyT firstKey = condSelectKey.apply(
            c, leftKey, rightKey, par);
        return () -> new Pair<>(firstKey, firstValue);
      }).par((par, data) -> {
        List<ValueT> lastValue = subtractionValue.add(xorValue, data.getSecond(), par);
        KeyT lastKey = subtractionKey.add(xorKey, data.getFirst(), par);
        List<Pair<KeyT, List<ValueT>>> result = new ArrayList<>();
        result.add(data);
        result.add(new Pair<>(lastKey, lastValue));
        return () -> result;
      });
    }
  
    private interface Addition<ValueT, BuilderT> {
  
      static <V, B> Addition<List<V>, B> forLists(Addition<V, B> op) {
        return (a, b, builder) -> {
          List<V> result = new ArrayList<>();
          for (int i = 0; i < a.size(); i++) {
            result.add(op.add(a.get(i), b.get(i), builder));
          }
          return result;
        };
      }
  
      ValueT add(ValueT a, ValueT b, BuilderT builder);
    }
  
    private interface Compare<ValueT, ResultT, BuilderT> {
  
      ResultT compare(ValueT a, ValueT b, BuilderT builder);
    }
  
    private interface ConditionalSelect<ConditionT, ValueT, BuilderT> {
  
      static <C, V, B extends ProtocolBuilderImpl<B>> ConditionalSelect<C, List<V>, B> forLists(
          ConditionalSelect<C, V, B> op) {
        return (c, a, b, builder) -> {
          List<V> result = new ArrayList<>();
          for (int i = 0; i < a.size(); i++) {
            result.add(op.apply(c, a.get(i), b.get(i), builder));
          }
          return result;
        };
      }
  
      ValueT apply(ConditionT condition, ValueT a, ValueT b, BuilderT builder);
    }
  }
  