content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you 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 org.apache.cassandra.service.accord.txn;
  
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.Comparator;
  import java.util.RandomAccess;
  import java.util.function.Supplier;
  import javax.annotation.Nullable;
  
  import com.google.common.collect.Lists;
  
  import accord.primitives.Ranges;
  import org.apache.cassandra.db.TypeSizes;
  import org.apache.cassandra.db.partitions.FilteredPartition;
  import org.apache.cassandra.db.partitions.PartitionIterator;
  import org.apache.cassandra.db.partitions.PartitionIterators;
  import org.apache.cassandra.db.rows.Row;
  import org.apache.cassandra.db.rows.UnfilteredRowIterator;
  import org.apache.cassandra.db.rows.UnfilteredRowIteratorSerializer;
  import org.apache.cassandra.db.rows.UnfilteredRowIterators;
  import org.apache.cassandra.io.util.DataInputPlus;
  import org.apache.cassandra.io.util.DataOutputPlus;
  import org.apache.cassandra.schema.Schema;
  import org.apache.cassandra.schema.TableId;
  import org.apache.cassandra.schema.TableMetadata;
  import org.apache.cassandra.service.accord.TokenRange;
  import org.apache.cassandra.service.accord.api.TokenKey;
  import org.apache.cassandra.service.accord.serializers.Version;
  import org.apache.cassandra.utils.ObjectSizes;
  
  import static com.google.common.base.Preconditions.checkState;
  import static org.apache.cassandra.db.SerializationHeader.StableHeaderSerializer.STABLE;
  import static org.apache.cassandra.db.rows.DeserializationHelper.Flag.FROM_REMOTE;
  import static org.apache.cassandra.utils.ObjectSizes.sizeOfReferenceArray;
  
  public class TxnDataRangeValue extends ArrayList<FilteredPartition> implements TxnDataValue, RandomAccess
  {
      private static final long EMPTY_SIZE = ObjectSizes.measure(new TxnDataRangeValue(0)) - sizeOfReferenceArray(0);
  
      public TxnDataRangeValue() {}
  
      public TxnDataRangeValue(int size)
      {
          super(size);
      }
  
      @Override
      public Kind kind()
      {
          return Kind.range;
      }
  
      @Override
      public TxnDataRangeValue merge(TxnDataValue other)
      {
          if (isEmpty())
              return (TxnDataRangeValue)other;
  
          TxnDataRangeValue otherRange = (TxnDataRangeValue)other;
          if (otherRange.isEmpty())
              return this;
  
          TableId tableId = tableId();
          for (FilteredPartition partition : otherRange)
              checkState(partition.metadata().id.equals(tableId), "All values should be for the same table");
  
          addAll(((TxnDataRangeValue)other));
          return this;
      }
  
      @Override
      public TxnDataValue without(Ranges ranges)
      {
          TokenRange range = range();
          if (range == null || !ranges.intersects(range))
              return this;
  
          if (ranges.contains(range))
              return null;
  
          TxnDataRangeValue result = new TxnDataRangeValue();
          for (FilteredPartition p : this)
          {
              if (!ranges.contains(new TokenKey(p.metadata().id, p.partitionKey().getToken())))
                  result.add(p);
          }
          return result;
      }
  
      @Override
      public long maxTimestamp()
      {
          long max = 0;
          for (FilteredPartition partition : this)
              max = Math.max(max, partition.maxTimestamp());
          return max;
      }
  
      Supplier<PartitionIterator> toPartitionIterator(boolean reversed)
      {
          // Sorting isn't preserved when merging TxnDataRangeValues together so sort here
          sort();
          return () -> PartitionIterators.concat(Lists.transform(this, v -> PartitionIterators.singletonIterator(v.rowIterator(reversed))));
      }
  
      private void sort()
      {
          sort(Comparator.comparing(FilteredPartition::partitionKey));
      }
  
      public @Nullable TableId tableId()
      {
          if (isEmpty())
              return null;
          return get(0).metadata().id;
      }
  
      public @Nullable TokenRange range()
      {
          if (isEmpty())
              return null;
          FilteredPartition first = get(0);
          FilteredPartition last = get(size() - 1);
          return TokenRange.create(new TokenKey(first.metadata().id, first.partitionKey().getToken()),
                                   new TokenKey(last.metadata().id, last.partitionKey().getToken()));
      }
  
  
  
      @Override
      public long estimatedSizeOnHeap()
      {
          long size = EMPTY_SIZE + sizeOfReferenceArray(size());
          for (FilteredPartition partition : this)
          {
              Row staticRow = partition.staticRow();
              if (staticRow != null)
                  size += staticRow.unsharedHeapSize();
              for (Row row : partition)
                  size += row.unsharedHeapSize();
          }
  
          // TODO: Include the other parts of FilteredPartition after we rebase to pull in BTreePartitionData?
          return size;
      }
  
      public static final TxnDataValueSerializer<TxnDataRangeValue> serializer = new TxnDataValueSerializer<>()
      {
          @Override
          public void serialize(TxnDataRangeValue value, DataOutputPlus out, Version version) throws IOException
          {
              out.writeUnsignedVInt32(value.size());
  
              if (value.isEmpty())
                  return;
  
              TableId.serializer.serialize(value.tableId(), out, version.messageVersion());
              for (FilteredPartition partition : value)
              {
                  try (UnfilteredRowIterator iterator = partition.unfilteredIterator())
                  {
                      UnfilteredRowIteratorSerializer.serializer.serialize(iterator, out, version.messageVersion(), partition.rowCount(), STABLE, null);
                  }
              }
          }
  
          @Override
          public TxnDataRangeValue deserialize(DataInputPlus in, Version version) throws IOException
          {
              int numPartitions = in.readUnsignedVInt32();
              TxnDataRangeValue value = new TxnDataRangeValue(numPartitions);
              if (numPartitions == 0)
                  return value;
              TableMetadata metadata = Schema.instance.getExistingTableMetadata(TableId.deserialize(in));
              for (int i = 0; i < numPartitions; i++)
              {
                  UnfilteredRowIteratorSerializer.Header header = UnfilteredRowIteratorSerializer.serializer.deserializeHeader(metadata, in, version.messageVersion(), FROM_REMOTE, STABLE, null);
                  try (UnfilteredRowIterator partition = UnfilteredRowIteratorSerializer.serializer.deserialize(in, version.messageVersion(), metadata, FROM_REMOTE, header))
                  {
                      value.add(new FilteredPartition(UnfilteredRowIterators.filter(partition, 0)));
                  }
              }
              return value;
          }
  
          @Override
          public long serializedSize(TxnDataRangeValue value, Version version)
          {
              long size = TypeSizes.sizeofUnsignedVInt(value.size());
              if (value.size() == 0)
                  return size;
              size += TableId.serializer.serializedSize(value.tableId(), version.messageVersion());
              for (FilteredPartition partition : value)
              {
                  try (UnfilteredRowIterator iterator = partition.unfilteredIterator())
                  {
                      size += UnfilteredRowIteratorSerializer.serializer.serializedSize(iterator, version.messageVersion(), partition.rowCount(), STABLE, null);
                  }
              }
              return size;
          }
      };
  }
  