xref: /external/guava/guava/src/com/google/common/collect/TableCollectors.java

/*
 * Copyright (C) 2009 The Guava Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.google.common.collect;

import static com.google.common.base.Preconditions.checkNotNull;

import com.google.common.annotations.GwtCompatible;
import com.google.common.collect.Tables.AbstractCell;
import java.util.ArrayList;
import java.util.List;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collector;
import org.checkerframework.checker.nullness.qual.Nullable;

/** Collectors utilities for {@code common.collect.Table} internals. */
@GwtCompatible
@ElementTypesAreNonnullByDefault
final class TableCollectors {

  static <T extends @Nullable Object, R, C, V>
      Collector<T, ?, ImmutableTable<R, C, V>> toImmutableTable(
          Function<? super T, ? extends R> rowFunction,
          Function<? super T, ? extends C> columnFunction,
          Function<? super T, ? extends V> valueFunction) {
    checkNotNull(rowFunction, "rowFunction");
    checkNotNull(columnFunction, "columnFunction");
    checkNotNull(valueFunction, "valueFunction");
    return Collector.of(
        (Supplier<ImmutableTable.Builder<R, C, V>>) ImmutableTable.Builder::new,
        (builder, t) ->
            builder.put(rowFunction.apply(t), columnFunction.apply(t), valueFunction.apply(t)),
        ImmutableTable.Builder::combine,
        ImmutableTable.Builder::build);
  }

  static <T extends @Nullable Object, R, C, V>
      Collector<T, ?, ImmutableTable<R, C, V>> toImmutableTable(
          Function<? super T, ? extends R> rowFunction,
          Function<? super T, ? extends C> columnFunction,
          Function<? super T, ? extends V> valueFunction,
          BinaryOperator<V> mergeFunction) {

    checkNotNull(rowFunction, "rowFunction");
    checkNotNull(columnFunction, "columnFunction");
    checkNotNull(valueFunction, "valueFunction");
    checkNotNull(mergeFunction, "mergeFunction");

    /*
     * No mutable Table exactly matches the insertion order behavior of ImmutableTable.Builder, but
     * the Builder can't efficiently support merging of duplicate values.  Getting around this
     * requires some work.
     */

    return Collector.of(
        () -> new ImmutableTableCollectorState<R, C, V>()
        /* GWT isn't currently playing nicely with constructor references? */ ,
        (state, input) ->
            state.put(
                rowFunction.apply(input),
                columnFunction.apply(input),
                valueFunction.apply(input),
                mergeFunction),
        (s1, s2) -> s1.combine(s2, mergeFunction),
        state -> state.toTable());
  }

  static <
          T extends @Nullable Object,
          R extends @Nullable Object,
          C extends @Nullable Object,
          V extends @Nullable Object,
          I extends Table<R, C, V>>
      Collector<T, ?, I> toTable(
          java.util.function.Function<? super T, ? extends R> rowFunction,
          java.util.function.Function<? super T, ? extends C> columnFunction,
          java.util.function.Function<? super T, ? extends V> valueFunction,
          java.util.function.Supplier<I> tableSupplier) {
    return toTable(
        rowFunction,
        columnFunction,
        valueFunction,
        (v1, v2) -> {
          throw new IllegalStateException("Conflicting values " + v1 + " and " + v2);
        },
        tableSupplier);
  }

  static <
          T extends @Nullable Object,
          R extends @Nullable Object,
          C extends @Nullable Object,
          V extends @Nullable Object,
          I extends Table<R, C, V>>
      Collector<T, ?, I> toTable(
          java.util.function.Function<? super T, ? extends R> rowFunction,
          java.util.function.Function<? super T, ? extends C> columnFunction,
          java.util.function.Function<? super T, ? extends V> valueFunction,
          BinaryOperator<V> mergeFunction,
          java.util.function.Supplier<I> tableSupplier) {
    checkNotNull(rowFunction);
    checkNotNull(columnFunction);
    checkNotNull(valueFunction);
    checkNotNull(mergeFunction);
    checkNotNull(tableSupplier);
    return Collector.of(
        tableSupplier,
        (table, input) ->
            mergeTables(
                table,
                rowFunction.apply(input),
                columnFunction.apply(input),
                valueFunction.apply(input),
                mergeFunction),
        (table1, table2) -> {
          for (Table.Cell<R, C, V> cell2 : table2.cellSet()) {
            mergeTables(
                table1, cell2.getRowKey(), cell2.getColumnKey(), cell2.getValue(), mergeFunction);
          }
          return table1;
        });
  }

  private static final class ImmutableTableCollectorState<R, C, V> {
    final List<MutableCell<R, C, V>> insertionOrder = new ArrayList<>();
    final Table<R, C, MutableCell<R, C, V>> table = HashBasedTable.create();

    void put(R row, C column, V value, BinaryOperator<V> merger) {
      MutableCell<R, C, V> oldCell = table.get(row, column);
      if (oldCell == null) {
        MutableCell<R, C, V> cell = new MutableCell<>(row, column, value);
        insertionOrder.add(cell);
        table.put(row, column, cell);
      } else {
        oldCell.merge(value, merger);
      }
    }

    ImmutableTableCollectorState<R, C, V> combine(
        ImmutableTableCollectorState<R, C, V> other, BinaryOperator<V> merger) {
      for (MutableCell<R, C, V> cell : other.insertionOrder) {
        put(cell.getRowKey(), cell.getColumnKey(), cell.getValue(), merger);
      }
      return this;
    }

    ImmutableTable<R, C, V> toTable() {
      return ImmutableTable.copyOf(insertionOrder);
    }
  }

  private static final class MutableCell<R, C, V> extends AbstractCell<R, C, V> {
    private final R row;
    private final C column;
    private V value;

    MutableCell(R row, C column, V value) {
      this.row = checkNotNull(row, "row");
      this.column = checkNotNull(column, "column");
      this.value = checkNotNull(value, "value");
    }

    @Override
    public R getRowKey() {
      return row;
    }

    @Override
    public C getColumnKey() {
      return column;
    }

    @Override
    public V getValue() {
      return value;
    }

    void merge(V value, BinaryOperator<V> mergeFunction) {
      checkNotNull(value, "value");
      this.value = checkNotNull(mergeFunction.apply(this.value, value), "mergeFunction.apply");
    }
  }

  private static <
          R extends @Nullable Object, C extends @Nullable Object, V extends @Nullable Object>
      void mergeTables(
          Table<R, C, V> table,
          @ParametricNullness R row,
          @ParametricNullness C column,
          @ParametricNullness V value,
          BinaryOperator<V> mergeFunction) {
    checkNotNull(value);
    V oldValue = table.get(row, column);
    if (oldValue == null) {
      table.put(row, column, value);
    } else {
      V newValue = mergeFunction.apply(oldValue, value);
      if (newValue == null) {
        table.remove(row, column);
      } else {
        table.put(row, column, newValue);
      }
    }
  }

  private TableCollectors() {}
}