java-commons/commons-serialize-databind/src/main/java/io/gitlab/jfronny/commons/serialize/databind/impl/adapter/MapTypeAdapterFactory.java

package io.gitlab.jfronny.commons.serialize.databind.impl.adapter;

import io.gitlab.jfronny.commons.serialize.databind.ObjectMapper;
import io.gitlab.jfronny.commons.serialize.databind.api.TypeAdapter;
import io.gitlab.jfronny.commons.serialize.databind.TypeAdapterFactory;
import io.gitlab.jfronny.commons.serialize.databind.api.TypeToken;
import io.gitlab.jfronny.commons.serialize.databind.impl.MapKeyReader;
import io.gitlab.jfronny.commons.serialize.databind.impl.MapKeyWriter;
import io.gitlab.jfronny.commons.serialize.databind.api.TypeUtils;
import io.gitlab.jfronny.commons.serialize.MalformedDataException;
import io.gitlab.jfronny.commons.serialize.SerializeReader;
import io.gitlab.jfronny.commons.serialize.SerializeWriter;
import io.gitlab.jfronny.commons.serialize.Token;
import io.gitlab.jfronny.commons.tuple.Tuple;

import java.lang.reflect.Type;
import java.util.*;

public class MapTypeAdapterFactory implements TypeAdapterFactory {
    @Override
    public <T> TypeAdapter<T> create(ObjectMapper mapper, TypeToken<T> typeToken) {
        Type type = typeToken.getType();

        Class<? super T> rawType = typeToken.getRawType();
        if (!Map.class.isAssignableFrom(rawType)) {
            return null;
        }

        Type[] keyAndValueTypes = TypeUtils.getMapKeyAndValueTypes(type, rawType);
        TypeAdapter<?> keyAdapter = mapper.getAdapter(TypeToken.get(keyAndValueTypes[0]));
        TypeAdapter<?> valueAdapter = mapper.getAdapter(TypeToken.get(keyAndValueTypes[1]));

        @SuppressWarnings({"unchecked", "rawtypes"})
        // we don't define a type parameter for the key or value types
        TypeAdapter<T> result = new MapTypeAdapter(mapper, keyAndValueTypes[0], keyAdapter, keyAndValueTypes[1], valueAdapter, rawType);
        return result;
    }

    private static class MapTypeAdapter<K, V> extends TypeAdapter<Map<K, V>> {
        private static final Set<Class<?>> RAW_TYPES_L = Set.of(Map.class, SequencedMap.class, AbstractMap.class);
        private static final Set<Class<?>> RAW_TYPES_S = Set.of(SortedMap.class, NavigableMap.class);

        private final TypeAdapter<K> keyTypeAdapter;
        private final TypeAdapter<V> valueTypeAdapter;
        private final Class<?> implClass;

        public MapTypeAdapter(
                ObjectMapper context,
                Type keyType,
                TypeAdapter<K> keyTypeAdapter,
                Type valueType,
                TypeAdapter<V> valueTypeAdapter,
                Class<?> implClass) {
            this.keyTypeAdapter = new TypeAdapterRuntimeTypeWrapper<>(context, keyTypeAdapter, keyType);
            this.valueTypeAdapter = new TypeAdapterRuntimeTypeWrapper<>(context, valueTypeAdapter, valueType);
            this.implClass = RAW_TYPES_L.contains(implClass) ? LinkedHashMap.class
                    : RAW_TYPES_S.contains(implClass) ? TreeMap.class
                    : implClass;
        }

        @Override
        public <TEx extends Exception, Writer extends SerializeWriter<TEx, Writer>> void serialize(Map<K, V> value, Writer writer) throws TEx, MalformedDataException {
            writer.beginObject();
            Map<String, Tuple<V, List<String>>> tmp = new LinkedHashMap<>();
            List<Map.Entry<K, V>> toWrite = new ArrayList<>(value.entrySet());
            int written = 0;
            for (Map.Entry<K, V> entry : toWrite) {
                try (MapKeyWriter keyWriter = new MapKeyWriter()) {
                    try {
                        keyTypeAdapter.serialize(entry.getKey(), keyWriter);
                    } catch (MalformedDataException e) {
                        if (keyWriter.isFailureOrigin()) {
                            // This mode of serialization won't work for the type in use. Try again with array serialization.
                            writer.beginArray();
                            // Write all entries that have been written so far
                            for (Map.Entry<String, Tuple<V, List<String>>> se : tmp.entrySet()) {
                                for (String comment : se.getValue().right()) {
                                    writer.comment(comment);
                                }
                                writer.name(se.getKey());
                                valueTypeAdapter.serialize(se.getValue().left(), writer);
                            }
                            // Write the rest of the entries
                            for (Map.Entry<K, V> se : toWrite) {
                                if (written-- > 0) continue;
                                writer.beginArray();
                                keyTypeAdapter.serialize(se.getKey(), writer);
                                valueTypeAdapter.serialize(se.getValue(), writer);
                                writer.endArray();
                            }
                            writer.endArray();
                            return;
                        }
                    }
                    tmp.put(keyWriter.getResult(), Tuple.of(entry.getValue(), keyWriter.getComments()));
                    written++;
                }
            }
            for (Map.Entry<String, Tuple<V, List<String>>> entry : tmp.entrySet()) {
                for (String comment : entry.getValue().right()) {
                    writer.comment(comment);
                }
                writer.name(entry.getKey());
                valueTypeAdapter.serialize(entry.getValue().left(), writer);
            }
            writer.endObject();
        }

        @Override
        public <TEx extends Exception, Reader extends SerializeReader<TEx, Reader>> Map<K, V> deserialize(Reader reader) throws TEx, MalformedDataException {
            Token peek = reader.peek();
            Map<K, V> map = (Map<K, V>) TypeUtils.instantiate(implClass).orElseThrow(() -> new MalformedDataException("Could not instantiate map"));
            if (peek == Token.BEGIN_ARRAY) {
                reader.beginArray();
                while (reader.hasNext()) {
                    reader.beginArray(); // entry array
                    K key = keyTypeAdapter.deserialize(reader);
                    V value = valueTypeAdapter.deserialize(reader);
                    if (map.put(key, value) != null && !reader.isLenient()) {
                        throw new MalformedDataException("duplicate key: " + key);
                    }
                    reader.endArray();
                }
                reader.endArray();
                return map;
            } else if (peek == Token.BEGIN_OBJECT) {
                reader.beginObject();
                while (reader.hasNext()) {
                    K key;
                    try (MapKeyReader keyReader = new MapKeyReader(reader.getPath(), reader.getPreviousPath(), reader.nextName())) {
                        key = keyTypeAdapter.deserialize(keyReader);
                    }
                    V value = valueTypeAdapter.deserialize(reader);
                    if (map.put(key, value) != null && !reader.isLenient()) {
                        throw new MalformedDataException("duplicate key: " + key);
                    }
                }
                reader.endObject();
                return map;
            } else {
                throw new MalformedDataException("Expected object or array but was " + peek);
            }
        }
    }
}