Polymorphic deserialization with Moshi

Example.java

package com.square.moshi.example;

import com.squareup.moshi.RuntimeTypeJsonAdapterFactory.RuntimeType;

public class Example {

    static class Animal {

        String type;
        String name;
    }

    static class Dog extends Animal {

        boolean playsCatch;
    }

    static class Cat extends Animal {

        boolean chasesRedLaserDot;
    }

    public static void main(String[] args) throws Exception {
        final RuntimeType typeInfo = RuntimeType.of(Animal.class, "type")
            .withSubtype(Dog.class, "dog")
            .withSubtype(Cat.class, "cat")
            .build();

        final RuntimeTypeJsonAdapterFactory factory = new RuntimeTypeJsonAdapterFactory()
            .registerRuntimeType(typeInfo);

        final Moshi moshi = new Moshi.Builder()
            .add(factory)
            .build();

        final JsonAdapter<Animal> animalJsonAdapter = moshi.adapter(Animal.class);
        final Dog dog = (Dog) animalJsonAdapter.fromJson("{\"type\":\"dog\",\"name\":\"Odie\",\"playsCatch\":\"true\"}");
        final Cat cat = (Cat) animalJsonAdapter.fromJson("{\"type\":\"cat\",\"name\":\"Garfield\",\"chasesRedLaserDot\":\"false\"}");
        animalJsonAdapter.toJson(dog);
        animalJsonAdapter.toJson(cat);
    }
}

RuntimeTypeJsonAdapterFactory.java

package com.squareup.moshi;

import java.io.EOFException;
import java.io.IOException;
import java.lang.annotation.Annotation;
import java.lang.reflect.Type;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Set;
import javax.annotation.Nullable;
import okio.Buffer;

/**
 * A {@link JsonAdapter.Factory} to handle polymorphic deserialization.
 */
public final class RuntimeTypeJsonAdapterFactory implements JsonAdapter.Factory {

    private final Map<Class<?>, RuntimeType> baseTypeToRuntimeType = new LinkedHashMap<>();

    public RuntimeTypeJsonAdapterFactory registerRuntimeType(RuntimeType type) {
        baseTypeToRuntimeType.put(type.baseType, type);
        return this;
    }

    @Nullable
    @Override
    @SuppressWarnings("unchecked")
    public JsonAdapter<?> create(Type type, Set<? extends Annotation> annotations, Moshi moshi) {
        if (!annotations.isEmpty()) {
            return null;
        }

        final RuntimeType runtimeType = baseTypeToRuntimeType.get(Types.getRawType(type));

        if (runtimeType == null) {
            return null;
        }

        final Map<String, JsonAdapter<Object>> discriminatorValueToJsonAdapter = new LinkedHashMap<>(runtimeType.discriminatorValueToSubtype.size());

        for (String key : runtimeType.discriminatorValueToSubtype.keySet()) {
            discriminatorValueToJsonAdapter.put(key, moshi.adapter(runtimeType.discriminatorValueToSubtype.get(key), annotations));
        }

        return new RuntimeTypeJsonAdapter(runtimeType, discriminatorValueToJsonAdapter);
    }

    /**
     * Encapsulates all the information needed to
     * identify a polymorphic deserialization case.
     */
    public static class RuntimeType {

        /**
         * The Json key whose value will determine
         * which type of object is.
         */
        public final String discriminatorKey;

        /**
         * The base class that every subtype extends.
         * This class will contain a field to store the
         * discriminator value, as well as some other
         * common properties.
         */
        public final Class<?> baseType;

        /**
         * A map that defines a subtype for each discriminator value.
         */
        public final Map<String, Class<?>> discriminatorValueToSubtype;

        private <T> RuntimeType(Builder<T> builder) {
            this.discriminatorKey = builder.discriminatorKey;
            this.baseType = builder.baseType;
            this.discriminatorValueToSubtype = Collections.unmodifiableMap(builder.discriminatorValueToSubtype);
        }

        public static <T> RuntimeType.Builder<T> of(Class<T> baseType, String discriminatorKey) {
            return new Builder<>(baseType, discriminatorKey);
        }

        public static class Builder<T> {

            private String discriminatorKey;
            private Class<T> baseType;
            private Map<String, Class<?>> discriminatorValueToSubtype = new LinkedHashMap<>();

            private Builder(Class<T> baseType, String discriminatorKey) {
                this.baseType = baseType;
                this.discriminatorKey = discriminatorKey;
            }

            /**
             * Stores the {{@code discriminatorValue}, {@code subtype}} pair each time
             * that gets called.
             */
            public Builder<T> withSubtype(Class<? extends T> subtype, String discriminatorValue) {
                discriminatorValueToSubtype.put(discriminatorValue, subtype);
                return this;
            }

            public RuntimeType build() {
                if (discriminatorKey == null) {
                    throw new IllegalArgumentException("discriminatorKey cannot be null");
                }

                if (discriminatorKey.isEmpty()) {
                    throw new IllegalArgumentException("discriminatorKey cannot be empty");
                }

                if (baseType == null) {
                    throw new IllegalArgumentException("baseType cannot be null");
                }

                return new RuntimeType(this);
            }
        }
    }

    private static class RuntimeTypeJsonAdapter extends JsonAdapter<Object> {

        private final RuntimeType runtimeType;
        private final JsonReader.Options options;
        private final Map<String, JsonAdapter<Object>> discriminatorValueToJsonAdapter;

        private RuntimeTypeJsonAdapter(RuntimeType runtimeType, Map<String, JsonAdapter<Object>> discriminatorValueToJsonAdapter) {
            this.runtimeType = runtimeType;
            this.options = JsonReader.Options.of(runtimeType.discriminatorKey);
            this.discriminatorValueToJsonAdapter = discriminatorValueToJsonAdapter;
        }

        @Nullable
        @Override
        public Object fromJson(JsonReader reader) throws IOException {
            // The idea of using the CopyAndReadJsonReader is to write
            // in a temp buffer everything that gets read from the original
            // reader until the "discriminatorKey" is found.
            final CopyAndReadJsonReader copyAndReadJsonReader = new CopyAndReadJsonReader(reader);

            copyAndReadJsonReader.beginObject();

            while (copyAndReadJsonReader.hasNext()) {
                switch (copyAndReadJsonReader.selectName(options)) {
                    case -1:
                        // Keep reading tokens until the "discriminatorKey" is found
                        copyAndReadJsonReader.nextToken();
                        break;
                    case 0:
                        // Note that the both "discriminatorKey" and "discriminatorValue" will
                        // be written into the temp buffer.
                        final String discriminatorValue = copyAndReadJsonReader.nextString();
                        final JsonAdapter<?> discriminatorAdapter = discriminatorValueToJsonAdapter.get(discriminatorValue);

                        if (discriminatorAdapter != null) {
                            // The idea of using the MergedJsonReader is that we read from the
                            // temp buffer until it gets exhausted and then switch back to
                            // the original JsonReader so we can keep parsing.
                            final JsonReader newReader = new MergedJsonReader(JsonReader.of(copyAndReadJsonReader.buffer), reader);
                            return discriminatorAdapter.fromJson(newReader);
                        }
                        break;
                }
            }

            copyAndReadJsonReader.endObject();

            return null;
        }

        @Override
        public void toJson(JsonWriter writer, @Nullable Object value) throws IOException {
            for (String key : runtimeType.discriminatorValueToSubtype.keySet()) {
                final Class<?> subtype = runtimeType.discriminatorValueToSubtype.get(key);
                Class<?> valueType = value.getClass();
                while (valueType != Object.class) {
                    if (valueType == subtype) {
                        discriminatorValueToJsonAdapter.get(key).toJson(value);
                        return;
                    }

                    valueType = valueType.getSuperclass();
                }
            }

            writer.nullValue();
        }
    }

    /**
     * A {@link JsonReader} that copies into a buffer everything
     * that gets read from the {@code delegate} {@link JsonReader}.
     */
    private static class CopyAndReadJsonReader extends JsonReader {

        private final JsonReader delegate;
        private final Buffer buffer;
        private final JsonWriter writer;

        private CopyAndReadJsonReader(JsonReader delegate) {
            this.delegate = delegate;
            this.buffer = new Buffer();
            this.writer = JsonWriter.of(buffer);
        }

        @Override public void beginArray() throws IOException {
            delegate.beginArray();
            writer.beginArray();
        }

        @Override public void endArray() throws IOException {
            delegate.endArray();
            writer.endArray();
        }

        @Override public void beginObject() throws IOException {
            delegate.beginObject();
            writer.beginObject();
        }

        @Override public void endObject() throws IOException {
            delegate.endObject();
            writer.endObject();
        }

        @Override public boolean hasNext() throws IOException {
            return delegate.hasNext();
        }

        @Override public Token peek() throws IOException {
            return delegate.peek();
        }

        @Override public String nextName() throws IOException {
            final String nextName = delegate.nextName();
            writer.name(nextName);
            return nextName;
        }

        @Override public int selectName(Options options) throws IOException {
            int result = delegate.selectName(options);
            if (result != -1) {
                writer.name(options.strings[result]);
            }
            return result;
        }

        @Override public String nextString() throws IOException {
            final String nextString = delegate.nextString();
            writer.value(nextString);
            return nextString;
        }

        @Override public int selectString(Options options) throws IOException {
            return delegate.selectString(options);
        }

        @Override public boolean nextBoolean() throws IOException {
            final boolean nextBoolean = delegate.nextBoolean();
            writer.value(nextBoolean);
            return nextBoolean;
        }

        @Nullable @Override public <T> T nextNull() throws IOException {
            writer.nullValue();
            return delegate.nextNull();
        }

        @Override public double nextDouble() throws IOException {
            final double nextDouble = delegate.nextDouble();
            writer.value(nextDouble);
            return nextDouble;
        }

        @Override public long nextLong() throws IOException {
            final long nextLong = delegate.nextLong();
            writer.value(nextLong);
            return nextLong;
        }

        @Override public int nextInt() throws IOException {
            final int nextInt = delegate.nextInt();
            writer.value(nextInt);
            return nextInt;
        }

        @Override public void skipValue() throws IOException {
            delegate.skipValue();
        }

        @Override void promoteNameToValue() throws IOException {
            writer.promoteValueToName();
            delegate.promoteNameToValue();
        }

        @Override public void close() throws IOException {
            delegate.close();
            writer.close();
        }

        public void nextToken() throws IOException {
            switch (peek()) {
                case BEGIN_ARRAY:
                    beginArray();
                    break;
                case END_ARRAY:
                    endArray();
                    break;
                case BEGIN_OBJECT:
                    beginObject();
                    break;
                case END_OBJECT:
                    endObject();
                    break;
                case NAME:
                    nextName();
                    break;
                case NUMBER:
                    try {
                        nextLong();
                    } catch (Exception ignored) {
                        nextDouble();
                    }
                    break;
                case BOOLEAN:
                    nextBoolean();
                    break;
                case STRING:
                    nextString();
                    break;
                case NULL:
                    nextNull();
                    break;
            }
        }
    }

    /**
     * A {@link JsonReader} that receives a list of {@link JsonReader},
     * starts reading from the first one and switches to next one
     * as soon as the current one gets exhausted.
     */
    private static class MergedJsonReader extends JsonReader {

        private final JsonReader[] readers;

        private JsonReader currentReader;
        private int currentReaderIndex = 0;

        private MergedJsonReader(JsonReader... jsonReaders) {
            this.readers = jsonReaders;
            this.currentReader = jsonReaders[0];
        }

        @Override public void beginArray() throws IOException {
            currentReader.beginArray();
        }

        @Override public void endArray() throws IOException {
            currentReader.endArray();
        }

        @Override public void beginObject() throws IOException {
            currentReader.beginObject();
        }

        @Override public void endObject() throws IOException {
            currentReader.endObject();
        }

        @Override public boolean hasNext() throws IOException {
            try {
                return currentReader.hasNext();
            } catch (EOFException e) {
                currentReaderIndex++;
                currentReader = readers[currentReaderIndex];
            }
            return hasNext();
        }

        @Override public Token peek() throws IOException {
            return currentReader.peek();
        }

        @Override public String nextName() throws IOException {
            return currentReader.nextName();
        }

        @Override public int selectName(Options options) throws IOException {
            return currentReader.selectName(options);
        }

        @Override public String nextString() throws IOException {
            return currentReader.nextString();
        }

        @Override public int selectString(Options options) throws IOException {
            return currentReader.selectString(options);
        }

        @Override public boolean nextBoolean() throws IOException {
            return currentReader.nextBoolean();
        }

        @Nullable @Override public <T> T nextNull() throws IOException {
            return currentReader.nextNull();
        }

        @Override public double nextDouble() throws IOException {
            return currentReader.nextDouble();
        }

        @Override public long nextLong() throws IOException {
            return currentReader.nextLong();
        }

        @Override public int nextInt() throws IOException {
            return currentReader.nextInt();
        }

        @Override public void skipValue() throws IOException {
            currentReader.skipValue();
        }

        @Override void promoteNameToValue() throws IOException {
            currentReader.promoteNameToValue();
        }

        @Override public void close() throws IOException {
            for (JsonReader reader : readers) {
                reader.close();
            }
        }
    }
}