GeneratorContext.java

package net.morimekta.providence.testing.generator;

import net.morimekta.providence.PMessage;
import net.morimekta.providence.PMessageOrBuilder;
import net.morimekta.providence.descriptor.PDescriptor;
import net.morimekta.providence.descriptor.PEnumDescriptor;
import net.morimekta.providence.descriptor.PList;
import net.morimekta.providence.descriptor.PMap;
import net.morimekta.providence.descriptor.PMessageDescriptor;
import net.morimekta.providence.descriptor.PSet;
import net.morimekta.providence.testing.generator.defaults.BinaryGenerator;
import net.morimekta.providence.testing.generator.defaults.BoolGenerator;
import net.morimekta.providence.testing.generator.defaults.ByteGenerator;
import net.morimekta.providence.testing.generator.defaults.DoubleGenerator;
import net.morimekta.providence.testing.generator.defaults.EnumGenerator;
import net.morimekta.providence.testing.generator.defaults.IntGenerator;
import net.morimekta.providence.testing.generator.defaults.ListGenerator;
import net.morimekta.providence.testing.generator.defaults.LongGenerator;
import net.morimekta.providence.testing.generator.defaults.MapGenerator;
import net.morimekta.providence.testing.generator.defaults.SetGenerator;
import net.morimekta.providence.testing.generator.defaults.ShortGenerator;
import net.morimekta.providence.testing.generator.defaults.StringGenerator;
import net.morimekta.providence.testing.generator.defaults.VoidGenerator;
import net.morimekta.util.collect.UnmodifiableList;

import javax.annotation.Nonnull;
import java.lang.reflect.Constructor;
import java.lang.reflect.InvocationTargetException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;
import java.util.function.Supplier;

/**
 * This is the core generator context, which essentially brings all things to be generated together.
 * It has full knowledge about all the generators, where state is more or less copied from
 * parent generators. And the current generating context.
 *
 * @param <Context> The context implementation type.
 */
public class GeneratorContext<Context extends GeneratorContext<Context>> {
    /**
     * Create a fresh generator.
     */
    public GeneratorContext() {
        this.random = new Random();
        this.generatorMap = new ConcurrentHashMap<>();
        this.persistentProperties = new ConcurrentHashMap<>();
        this.fillRate = 1.0;
        this.minCollectionSize = 0;
        this.naxCollectionSize = 10;
        this.generatedMessages = Collections.synchronizedList(new ArrayList<>());
    }

    /**
     * Create a generator instance as copy of another instance. The
     * new generator can then be updated without interfering with the
     * parent.
     *
     * @param parent The parent generator instance.
     */
    protected GeneratorContext(GeneratorContext<Context> parent) {
        this.random = parent.random;
        this.generatorMap = new ConcurrentHashMap<>(parent.generatorMap);
        this.persistentProperties = new ConcurrentHashMap<>(parent.persistentProperties);
        this.fillRate = parent.fillRate;
        this.minCollectionSize = parent.minCollectionSize;
        this.naxCollectionSize = parent.naxCollectionSize;
        this.generatedMessages = parent.generatedMessages;
    }

    @SuppressWarnings("unchecked")
    public Context deepCopy() {
        try {
            Constructor constructor = getClass().getConstructor(getClass());
            return (Context) constructor.newInstance(this);
        } catch (NoSuchMethodException | InvocationTargetException | InstantiationException | IllegalAccessException e) {
            throw new AssertionError("No usable copy constructor for " + getClass(), e);
        }
    }

    public Context setFillRate(double v) {
        fillRate = v;
        return self();
    }

    public Context setMinCollectionSize(int v) {
        minCollectionSize = v;
        return self();
    }

    public Context setMaxCollectionSize(int v) {
        naxCollectionSize = v;
        return self();
    }

    public Random getRandom() {
        return random;
    }

    public Context setRandom(Random random) {
        this.random = random;
        return self();
    }

    @SuppressWarnings("unchecked")
    public <V> V getProperty(String name) {
        return (V) persistentProperties.get(name);
    }

    @SuppressWarnings("unchecked")
    public <V> V createPropertyIfAbsent(String name, Supplier<V> propertySupplier) {
        return (V) persistentProperties.computeIfAbsent(name, n -> propertySupplier.get());
    }

    public <V> Context setProperty(String name, V value) {
        persistentProperties.put(name, value);
        return self();
    }

    /**
     * Set a specific value generator based on type for default value generator.
     *
     * @param descriptor The type to be generated.
     * @param generator The value generator.
     * @param <V> The value type.
     * @return The generator instance.
     */
    @Nonnull
    public <V> Context withGenerator(PDescriptor descriptor, Generator<Context, V> generator) {
        generatorMap.put(descriptor, generator);
        return self();
    }

    @Nonnull
    @SuppressWarnings("unchecked")
    public <M extends PMessage<M>, VG extends Generator<Context, M>>
    VG generatorFor(@Nonnull PMessageDescriptor<M> descriptor) {
        return (VG) generatorForDescriptor(descriptor);
    }

    @Nonnull
    public Generator<Context, ?> generatorForDescriptor(@Nonnull PDescriptor descriptor) {
        return generatorMap.computeIfAbsent(descriptor, (d) -> makeGeneratorInternal(descriptor));
    }

    /**
     * Update the generator for the given message type.
     *
     * @param descriptor The message descriptor.
     * @param closure Closure updating the generator.
     * @param <M> The message type to have updated generator.
     * @param <MOB> The message or builder interface.
     * @return The generator instance.
     */
    @Nonnull
    @SuppressWarnings("unchecked")
    public <M extends PMessage<M>, MOB extends PMessageOrBuilder<M>>
    GeneratorContext<Context> withMessageGenerator(@Nonnull PMessageDescriptor<M> descriptor,
                                                   @Nonnull Consumer<MessageGenerator<Context, M, MOB>> closure) {
        generatorMap.compute(descriptor, (d, gen) -> {
            MessageGenerator<Context, M, MOB> mg = (MessageGenerator<Context, M, MOB>) gen;
            if (mg == null) {
                mg = new MessageGenerator<>(descriptor);
            } else {
                mg = mg.deepCopy();
            }
            closure.accept(mg);
            return mg;
        });
        return this;
    }

    /**
     * Create a message based on the current generator.
     *
     * @param descriptor The message descriptor.
     * @param <M> The message type.
     * @return The generated message.
     */
    @SuppressWarnings("unchecked")
    public <M extends PMessage<M>>
    M nextMessage(@Nonnull PMessageDescriptor<M> descriptor) {
        return (M) generatorMap.computeIfAbsent(descriptor, (d) -> new MessageGenerator<>(descriptor))
                               .generate(self());
    }

    /**
     * Create a message based on a modified generator.
     *
     * @param descriptor The message descriptor.
     * @param closure Closure updating the generator for this case only.
     * @param <M> The message type.
     * @param <MOB> The message or builder interface.
     * @return The generated message.
     */
    @SuppressWarnings("unchecked")
    public <M extends PMessage<M>, MOB extends PMessageOrBuilder<M>>
    M nextMessage(@Nonnull PMessageDescriptor<M> descriptor,
                  @Nonnull Consumer<MessageGenerator<Context, M, MOB>> closure) {
        MessageGenerator<Context, M, MOB> generator = (MessageGenerator<Context, M, MOB>) generatorMap.get(descriptor);
        if (generator == null) {
            generator = new MessageGenerator<>(descriptor);
        } else {
            generator = generator.deepCopy();
        }
        closure.accept(generator);
        return generator.generate(self());
    }

    public int nextCollectionSize() {
        return minCollectionSize +
               getRandom().nextInt(naxCollectionSize - minCollectionSize);
    }

    public boolean nextFieldIsPresent() {
        return getRandom().nextDouble() < fillRate;
    }


    public List<PMessage> getGeneratedMessages() {
        return UnmodifiableList.copyOf(generatedMessages);
    }

    public void clearGeneratedMessages() {
        generatedMessages.clear();
    }

    // -------------------------
    // ----    PROTECTED    ----
    // -------------------------

    void addGeneratedMessage(PMessage message) {
        generatedMessages.add(message);
    }

    // -----------------------
    // ----    PRIVATE    ----
    // -----------------------

    private final Map<PDescriptor, Generator<Context, ?>> generatorMap;
    private final Map<String, Object>                     persistentProperties;
    private final List<PMessage>                          generatedMessages;

    private Random random;
    private double fillRate;
    private int    minCollectionSize;
    private int    naxCollectionSize;

    @Nonnull
    @SuppressWarnings("unchecked")
    private Context self() {
        return (Context) this;
    }

    @Nonnull
    @SuppressWarnings("unchecked")
    private <T> Generator<Context, T> makeGeneratorInternal(@Nonnull PDescriptor type) {
        switch (type.getType()) {
            case VOID:
                return new VoidGenerator();
            case BOOL:
                return new BoolGenerator();
            case BYTE:
                return new ByteGenerator();
            case I16:
                return new ShortGenerator();
            case I32:
                return new IntGenerator();
            case I64:
                return new LongGenerator();
            case DOUBLE:
                return new DoubleGenerator();
            case STRING:
                return new StringGenerator();
            case BINARY:
                return new BinaryGenerator();
            case LIST:
                return new ListGenerator((PList<Object>) type);
            case SET:
                return new SetGenerator((PSet<Object>) type);
            case MAP:
                return new MapGenerator((PMap<Object,Object>) type);
            case ENUM:
                return new EnumGenerator<>((PEnumDescriptor) type);
            case MESSAGE:
                return new MessageGenerator<>((PMessageDescriptor) type);
        }
        throw new IllegalArgumentException("Unhandled default type: " + type.getType());
    }

    /**
     * GeneratorContext with no extra methods in non-generic form.
     */
    public static final class Simple
            extends GeneratorContext<Simple> {
        public Simple() {
            super();
        }

        @SuppressWarnings("unused")
        public Simple(Simple generator) {
            super(generator);
        }
    }
}