Int2ObjectMapCodec.java

package com.hypixel.hytale.codec.codecs.map;

import com.hypixel.hytale.codec.Codec;
import com.hypixel.hytale.codec.ExtraInfo;
import com.hypixel.hytale.codec.WrappedCodec;
import com.hypixel.hytale.codec.codecs.StringIntegerCodec;
import com.hypixel.hytale.codec.exception.CodecException;
import com.hypixel.hytale.codec.schema.SchemaContext;
import com.hypixel.hytale.codec.schema.config.ObjectSchema;
import com.hypixel.hytale.codec.schema.config.Schema;
import com.hypixel.hytale.codec.schema.config.StringSchema;
import com.hypixel.hytale.codec.util.RawJsonReader;
import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
import it.unimi.dsi.fastutil.ints.Int2ObjectMaps;
import java.io.IOException;
import java.util.Map;
import java.util.function.Supplier;
import javax.annotation.Nonnull;
import org.bson.BsonDocument;
import org.bson.BsonValue;

public class Int2ObjectMapCodec<T> implements Codec<Int2ObjectMap<T>>, WrappedCodec<T> {
   private final Codec<T> valueCodec;
   private final Supplier<Int2ObjectMap<T>> supplier;
   private final boolean unmodifiable;

   public Int2ObjectMapCodec(Codec<T> valueCodec, Supplier<Int2ObjectMap<T>> supplier, boolean unmodifiable) {
      this.valueCodec = valueCodec;
      this.supplier = supplier;
      this.unmodifiable = unmodifiable;
   }

   public Int2ObjectMapCodec(Codec<T> valueCodec, Supplier<Int2ObjectMap<T>> supplier) {
      this(valueCodec, supplier, true);
   }

   public Codec<T> getChildCodec() {
      return this.valueCodec;
   }

   public Int2ObjectMap<T> decode(@Nonnull BsonValue bsonValue, @Nonnull ExtraInfo extraInfo) {
      BsonDocument bsonDocument = bsonValue.asDocument();
      Int2ObjectMap<T> map = (Int2ObjectMap)this.supplier.get();

      for(Map.Entry<String, BsonValue> entry : bsonDocument.entrySet()) {
         String key = (String)entry.getKey();
         BsonValue value = (BsonValue)entry.getValue();
         extraInfo.pushKey(key);

         try {
            int decodedKey = Integer.parseInt(key);
            map.put(decodedKey, this.valueCodec.decode(value, extraInfo));
         } catch (Exception e) {
            throw new CodecException("Failed to decode", value, extraInfo, e);
         } finally {
            extraInfo.popKey();
         }
      }

      if (this.unmodifiable) {
         map = Int2ObjectMaps.<T>unmodifiable(map);
      }

      return map;
   }

   @Nonnull
   public BsonValue encode(@Nonnull Int2ObjectMap<T> map, ExtraInfo extraInfo) {
      BsonDocument bsonDocument = new BsonDocument();

      for(Int2ObjectMap.Entry<T> entry : map.int2ObjectEntrySet()) {
         bsonDocument.put(Integer.toString(entry.getIntKey()), this.valueCodec.encode(entry.getValue(), extraInfo));
      }

      return bsonDocument;
   }

   public Int2ObjectMap<T> decodeJson(@Nonnull RawJsonReader reader, @Nonnull ExtraInfo extraInfo) throws IOException {
      reader.expect('{');
      reader.consumeWhiteSpace();
      Int2ObjectMap<T> map = (Int2ObjectMap)this.supplier.get();
      if (reader.tryConsume('}')) {
         if (this.unmodifiable) {
            map = Int2ObjectMaps.<T>unmodifiable(map);
         }

         return map;
      } else {
         while(true) {
            reader.expect('"');
            int decodedKey = reader.readIntValue();
            reader.expect('"');
            reader.consumeWhiteSpace();
            reader.expect(':');
            reader.consumeWhiteSpace();
            extraInfo.pushIntKey(decodedKey, reader);

            try {
               map.put(decodedKey, this.valueCodec.decodeJson(reader, extraInfo));
            } catch (Exception e) {
               throw new CodecException("Failed to decode", reader, extraInfo, e);
            } finally {
               extraInfo.popKey();
            }

            reader.consumeWhiteSpace();
            if (reader.tryConsumeOrExpect('}', ',')) {
               if (this.unmodifiable) {
                  map = Int2ObjectMaps.<T>unmodifiable(map);
               }

               return map;
            }

            reader.consumeWhiteSpace();
         }
      }
   }

   @Nonnull
   public Schema toSchema(@Nonnull SchemaContext context) {
      ObjectSchema s = new ObjectSchema();
      StringSchema name = StringIntegerCodec.INSTANCE.toSchema(context);
      s.setPropertyNames(name);
      s.setAdditionalProperties(context.refDefinition(this.valueCodec));
      return s;
   }
}