/*

  * 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.writable.kryo.serializers;

 import java.io.IOException;

 import java.io.ObjectInputStream;

 import java.io.ObjectOutputStream;

 import java.lang.reflect.InvocationTargetException;

 import java.lang.reflect.Method;

 import java.util.ArrayList;

 import com.esotericsoftware.kryo.Kryo;

 import com.esotericsoftware.kryo.Serializer;

 import com.esotericsoftware.kryo.io.Input;

 import com.esotericsoftware.kryo.io.Output;

 import com.esotericsoftware.kryo.serializers.FieldSerializer;

 /**

  * Kryo Serializer for Fastutil collection class.

  * By default, because they extend boxed collections, are being serialized very

  * inefficiently through a lot of temporary object creation.

  *

  * We are relying that fastutil classes are written to be correctly serialized

  * with Java serialization, and have put transient on all array fields and are

  * doing custom efficient serialization in writeObject/readObject methods.

  * This Serializer then swaps ObjectOutputStream for all default fields with

  * FieldSerializer, and then calls appropriate writeObject/readObject methods.

  * We are also relying on defaultWriteObject/defaultReadObject being

  * effectively called first within those methods

  *

  * @param <T> Object type

  */

 public class FastUtilSerializer<T> extends Serializer<T> {

   /** List of all types generated by fastutil */

   private static final String[] PRIMITIVE_TYPES = new String[] {

     "Boolean", "Byte", "Short", "Int", "Long", "Float", "Double", "Char",

     "Object"};

   /** List of all types used as keys in fastutil */

   private static final String[] PRIMITIVE_KEY_TYPES = new String[] {

     "Byte", "Short", "Int", "Long", "Float", "Double", "Char", "Object"};

   /** Field serializer for this fastutil class */

   private final FieldSerializer<T> fieldSerializer;

   /** Handle to writeObject Method on this fastutil class*/

   private final Method writeMethod;

   /** Handle to readObject Method on this fastutil class*/

   private final Method readMethod;

   /** Reusable output stream wrapper */

   private final FastUtilSerializer.FastutilKryoObjectOutputStream outputWrapper;

   /** Reusable input stream wrapper */

   private final FastUtilSerializer.FastutilKryoObjectInputStream inputWrapper;

   /**

    * Creates and initializes new serializer for a given fastutil class.

    * @param kryo Kryo instance

    * @param type Fastutil class

    */

   public FastUtilSerializer(Kryo kryo, Class<T> type) {

     fieldSerializer = new FieldSerializer<>(kryo, type);

     fieldSerializer.setIgnoreSyntheticFields(false);

     try {

       writeMethod = type.getDeclaredMethod(

           "writeObject", ObjectOutputStream.class);

       writeMethod.setAccessible(true);

       readMethod = type.getDeclaredMethod(

           "readObject", ObjectInputStream.class);

       readMethod.setAccessible(true);

     } catch (NoSuchMethodException e) {

       throw new RuntimeException(

           "Fastutil class " + type +

           " doesn't have readObject/writeObject methods", e);

     }

     try {

       outputWrapper = new FastutilKryoObjectOutputStream();

       inputWrapper = new FastutilKryoObjectInputStream();

     } catch (IOException e) {

       throw new RuntimeException(e);

     }

   }

   /**

    * Register serializer for a given fastutil class.

    * @param kryo Kryo instance

    * @param fastutilClass Fastutil class

    */

   public static void register(Kryo kryo, Class<?> fastutilClass) {

     kryo.register(fastutilClass, new FastUtilSerializer<>(kryo, fastutilClass));

   }

   /**

    * Registers serializers for all possible fastutil classes.

    *

    * There are many fastutil classes, so it is recommended to call this

    * function at the end, so they fastutil classes don't use up small

    * registration numbers.

    *

    * @param kryo Kryo instance

    */

   public static void registerAll(Kryo kryo) {

     registerArrayLists(kryo);

     registerArrayBigList(kryo);

     registerOpenHashSets(kryo);

     registerArraySets(kryo);

     registerRBTreeSets(kryo);

     registerAVLTreeSets(kryo);

     registerOpenHashMaps(kryo);

     registerRBTreeMaps(kryo);

     registerAVLTreeMaps(kryo);

     // Note - HeapPriorityQueues don't extend boxed collection,

     // and so they work out of the box correctly

   }

   /**

    * Register all Fastutil ArrayLists.

    *

    * @param kryo Kryo instance

    */

   public static void registerArrayLists(Kryo kryo) {

     registerAll(kryo, singleTypes(

         "it.unimi.dsi.fastutil._t1_s._T1_ArrayList", PRIMITIVE_TYPES));

   }

   /**

    * Register all Fastutil ArrayBigLists.

    *

    * @param kryo Kryo instance

    */

   public static void registerArrayBigList(Kryo kryo) {

     registerAll(kryo, singleTypes(

         "it.unimi.dsi.fastutil._t1_s._T1_BigArrayBigList", PRIMITIVE_TYPES));

   }

   /**

    * Register all Fastutil OpenHashSets.

    *

    * @param kryo Kryo instance

    */

   public static void registerOpenHashSets(Kryo kryo) {

     registerAll(kryo, singleTypes(

         "it.unimi.dsi.fastutil._t1_s._T1_OpenHashSet", PRIMITIVE_TYPES));

   }

   /**

    * Register all Fastutil ArraySets.

    *

    * @param kryo Kryo instance

    */

   public static void registerArraySets(Kryo kryo) {

     registerAll(kryo, singleTypes(

         "it.unimi.dsi.fastutil._t1_s._T1_ArraySet", PRIMITIVE_TYPES));

   }

   /**

    * Register all Fastutil RBTreeSets.

    *

    * @param kryo Kryo instance

    */

   public static void registerRBTreeSets(Kryo kryo) {

     registerAll(kryo, singleTypes(

         "it.unimi.dsi.fastutil._t1_s._T1_RBTreeSet", PRIMITIVE_KEY_TYPES));

   }

   /**

    * Register all Fastutil AVLTreeSets.

    *

    * @param kryo Kryo instance

    */

   public static void registerAVLTreeSets(Kryo kryo) {

     registerAll(kryo, singleTypes(

         "it.unimi.dsi.fastutil._t1_s._T1_AVLTreeSet", PRIMITIVE_KEY_TYPES));

   }

   /**

    * Register all Fastutil OpenHashMaps.

    *

    * @param kryo Kryo instance

    */

   public static void registerOpenHashMaps(Kryo kryo) {

     registerAll(kryo, doubleTypes(

         "it.unimi.dsi.fastutil._t1_s._T1_2_T2_OpenHashMap",

         PRIMITIVE_KEY_TYPES, PRIMITIVE_TYPES));

   }

   /**

    * Register all Fastutil RBTreeMaps.

    *

    * @param kryo Kryo instance

    */

   public static void registerRBTreeMaps(Kryo kryo) {

     registerAll(kryo, doubleTypes(

         "it.unimi.dsi.fastutil._t1_s._T1_2_T2_RBTreeMap",

         PRIMITIVE_KEY_TYPES, PRIMITIVE_TYPES));

   }

   /**

    * Register all Fastutil AVLTreeMaps.

    *

    * @param kryo Kryo instance

    */

   public static void registerAVLTreeMaps(Kryo kryo) {

     registerAll(kryo, doubleTypes(

         "it.unimi.dsi.fastutil._t1_s._T1_2_T2_AVLTreeMap",

         PRIMITIVE_KEY_TYPES, PRIMITIVE_TYPES));

   }

   /**

    * Register all class from the list of classes.

    *

    * @param kryo Kryo instance

    * @param types List of classes

    */

   private static void registerAll(Kryo kryo, ArrayList<Class<?>> types) {

     for (Class<?> type : types) {

       register(kryo, type);

     }

   }

   /**

    * Returns list of all classes that are generated by using given

    * pattern, and replacing it with passed list of types.

    * Pattern contains _t1_ and _T1_, for lowercase and actual name.

    *

    * @param pattern Given pattern

    * @param types Given list of strings to replace into pattern

    * @return List of all classes

    */

   private static ArrayList<Class<?>> singleTypes(

       String pattern, String[] types) {

     ArrayList<Class<?>> result = new ArrayList<>();

     for (String type : types) {

       try {

         result.add(Class.forName(

             pattern.replaceAll("_T1_", type).replaceAll(

                 "_t1_", type.toLowerCase())));

       } catch (ClassNotFoundException e) {

         throw new RuntimeException(pattern + " " + type, e);

       }

     }

     return result;

   }

   /**

    * Returns list of all classes that are generated by using given

    * pattern, and replacing it with passed list of types.

    * Pattern contains two variable pairs: _t1_, _T1_ and _t2_, _T2_,

    * in each pair one for lowercase and one for actual name.

    *

    * @param pattern Given pattern

    * @param types1 Given list of strings to replace t1 into pattern

    * @param types2 Given list of strings to replace t2 into pattern

    * @return List of all classes

    */

   private static ArrayList<Class<?>> doubleTypes(

       String pattern, String[] types1, String[] types2) {

     ArrayList<Class<?>> result = new ArrayList<>();

     for (String type1 : types1) {

       for (String type2 : types2) {

         try {

           result.add(Class.forName(

               pattern.replaceAll("_T1_", type1).replaceAll(

                   "_t1_", type1.toLowerCase())

                 .replaceAll("_T2_", type2).replaceAll(

                     "_t2_", type2.toLowerCase())));

         } catch (ClassNotFoundException e) {

           throw new RuntimeException(pattern + " " + type1 + " " + type2, e);

         }

       }

     }

     return result;

   }

   @Override

   public void write(Kryo kryo, Output output, T object) {

     fieldSerializer.write(kryo, output, object);

     outputWrapper.set(output, kryo);

     try {

       writeMethod.invoke(object, outputWrapper);

     } catch (IllegalAccessException | InvocationTargetException e) {

       throw new RuntimeException("writeObject failed", e);

     }

   }

   @Override

   public T read(Kryo kryo, Input input, Class<T> type) {

     T result = fieldSerializer.read(kryo, input, type);

     if (result != null) {

       inputWrapper.set(input, kryo);

       try {

         readMethod.invoke(result, inputWrapper);

       } catch (IllegalAccessException | InvocationTargetException e) {

         throw new RuntimeException("readObject failed", e);

       }

     }

     return result;

   }

   /**

    * Wrapper around ObjectOutputStream that ignores defaultWriteObject (assumes

    * that needed logic was already executed before), and passes all other calls

    * to Output

    */

   private static class FastutilKryoObjectOutputStream

       extends ObjectOutputStream {

     /** Output */

     private Output output;

     /** Kryo */

     private Kryo kryo;

     /** Constructor */

     FastutilKryoObjectOutputStream() throws IOException {

       super();

     }

     /**

      * Setter

      *

      * @param output Output

      * @param kryo kryo

      */

     public void set(Output output, Kryo kryo) {

       this.output = output;

       this.kryo = kryo;

     }

     @Override

     public void defaultWriteObject() throws IOException {

     }

     @Override

     public void writeBoolean(boolean val) throws IOException {

       output.writeBoolean(val);

     }

     @Override

     public void writeByte(int val) throws IOException  {

       output.writeByte(val);

     }

     @Override

     public void writeShort(int val)  throws IOException {

       output.writeShort(val);

     }

     @Override

     public void writeChar(int val)  throws IOException {

       output.writeChar((char) val);

     }

     @Override

     public void writeInt(int val)  throws IOException {

       output.writeInt(val, false);

     }

     @Override

     public void writeLong(long val)  throws IOException {

       output.writeLong(val, false);

     }

     @Override

     public void writeFloat(float val) throws IOException {

       output.writeFloat(val);

     }

     @Override

     public void writeDouble(double val) throws IOException {

       output.writeDouble(val);

     }

     @Override

     protected void writeObjectOverride(Object obj) throws IOException {

       kryo.writeClassAndObject(output, obj);

     }

   }

   /**

    * Wrapper around ObjectOutputStream that ignores defaultReadObject

    * (assumes that needed logic was already executed before), and passes

    * all other calls to Output

    */

   private static class FastutilKryoObjectInputStream extends ObjectInputStream {

     /** Input */

     private Input input;

     /** Kryo */

     private Kryo kryo;

     /** Constructor */

     FastutilKryoObjectInputStream() throws IOException {

       super();

     }

     /**

      * Setter

      *

      * @param input Input

      * @param kryo Kryo

      */

     public void set(Input input, Kryo kryo) {

       this.input = input;

       this.kryo = kryo;

     }

     @Override

     public void defaultReadObject() throws IOException, ClassNotFoundException {

     }

     @Override

     public boolean readBoolean() throws IOException {

       return input.readBoolean();

     }

     @Override

     public byte readByte() throws IOException {

       return input.readByte();

     }

     @Override

     public char readChar() throws IOException {

       return input.readChar();

     }

     @Override

     public short readShort() throws IOException {

       return input.readShort();

     }

     @Override

     public int readInt() throws IOException {

       return input.readInt(false);

     }

     @Override

     public long readLong() throws IOException {

       return input.readLong(false);

     }

     @Override

     public float readFloat() throws IOException {

       return input.readFloat();

     }

     @Override

     public double readDouble() throws IOException {

       return input.readDouble();

     }

     @Override

     protected Object readObjectOverride()

         throws IOException, ClassNotFoundException {

       return kryo.readClassAndObject(input);

     }

   }

 }