/*

  * 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.giraph.mapping;

 import it.unimi.dsi.fastutil.longs.Long2ObjectOpenHashMap;

 import java.util.Arrays;

 import java.util.Map;

 import java.util.concurrent.ConcurrentMap;

 import java.util.concurrent.atomic.AtomicLong;

 import javax.annotation.concurrent.ThreadSafe;

 import org.apache.giraph.conf.DefaultImmutableClassesGiraphConfigurable;

 import org.apache.giraph.conf.GiraphConstants;

 import org.apache.hadoop.io.ByteWritable;

 import org.apache.hadoop.io.LongWritable;

 import org.apache.hadoop.io.Writable;

 import org.apache.log4j.Logger;

 import com.google.common.collect.MapMaker;

 /**

  *

  * An implementation of MappingStore<LongWritable, ByteWritable>

  *

  * Methods implemented here are thread safe by default because it is guaranteed

  * that each entry is written to only once.

  * It can represent up to a maximum of 254 workers

  * any byte passed is treated as unsigned

  */

 @ThreadSafe

 public class LongByteMappingStore

   extends DefaultImmutableClassesGiraphConfigurable<LongWritable, Writable,

   Writable> implements MappingStore<LongWritable, ByteWritable> {

   /** Logger instance */

   private static final Logger LOG = Logger.getLogger(

     LongByteMappingStore.class);

   /** Counts number of entries added */

   private final AtomicLong numEntries = new AtomicLong(0);

   /** Id prefix to bytesArray index mapping */

   private ConcurrentMap<Long, byte[]> concurrentIdToBytes;

   /** Primitive idToBytes for faster querying */

   private Long2ObjectOpenHashMap<byte[]> idToBytes;

   /** Number of lower order bits */

   private int lower;

   /** Number of distinct prefixes */

   private int upper;

   /** Bit mask for lowerOrder suffix bits */

   private int lowerBitMask;

   /** LowerOrder bits count */

   private int lowerOrder;

   @Override

   public void initialize() {

     upper = GiraphConstants.LB_MAPPINGSTORE_UPPER.get(getConf());

     lower = GiraphConstants.LB_MAPPINGSTORE_LOWER.get(getConf());

     if ((lower & (lower - 1)) != 0) {

       throw new IllegalStateException("lower not a power of two");

     }

     lowerBitMask = lower - 1;

     lowerOrder = Integer.numberOfTrailingZeros(lower); // log_2_(lower)

     concurrentIdToBytes = new MapMaker()

         .initialCapacity(upper)

         .concurrencyLevel(getConf().getNumInputSplitsThreads())

         .makeMap();

     idToBytes = new Long2ObjectOpenHashMap<>(upper);

   }

   /**

    * Auxiliary method to be used by getTarget

    *

    * @param vertexId vertexId

    * @return return byte value of target

    */

   public byte getByteTarget(LongWritable vertexId) {

     long key = vertexId.get() >>> lowerOrder;

     int suffix = (int) (vertexId.get() & lowerBitMask);

     if (!idToBytes.containsKey(key)) {

       return -1;

     }

     return idToBytes.get(key)[suffix];

   }

   @Override

   public void addEntry(LongWritable vertexId, ByteWritable target) {

     long key = vertexId.get() >>> lowerOrder;

     byte[] bytes = concurrentIdToBytes.get(key);

     if (bytes == null) {

       byte[] newBytes = new byte[lower];

       Arrays.fill(newBytes, (byte) -1);

       bytes = concurrentIdToBytes.putIfAbsent(key, newBytes);

       if (bytes == null) {

         bytes = newBytes;

       }

     }

     bytes[(int) (vertexId.get() & lowerBitMask)] = target.get();

     numEntries.getAndIncrement(); // increment count

   }

   @Override

   public ByteWritable getTarget(LongWritable vertexId,

     ByteWritable target) {

     byte bval = getByteTarget(vertexId);

     if (bval == -1) { // worker not assigned by mapping

       return null;

     }

     target.set(bval);

     return target;

   }

   @Override

   public void postFilling() {

     // not thread-safe

     for (Map.Entry<Long, byte[]> entry : concurrentIdToBytes.entrySet()) {

       idToBytes.put(entry.getKey(), entry.getValue());

     }

     concurrentIdToBytes.clear();

     concurrentIdToBytes = null;

   }

   @Override

   public long getStats() {

     return numEntries.longValue();

   }

 }