<|start_filename|>src/main/java/com/medallia/data/DataSet.java<|end_filename|>
package com.medallia.data;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.function.Function;
import java.util.stream.Collectors;

public class DataSet {
	private final Map<String, FieldDefinition> fieldsByName;
	private List<FieldDefinition> fields;
	private List<Segment> segments;

	private DataSet(List<FieldSpec> fields, List<Segment> segments) {
		this.fields = new ArrayList<>();
		for (int col = 0; col < fields.size(); col++) {
			this.fields.add(new FieldDefinition(col, fields.get(col)));
		}

		this.segments = segments;
		this.fieldsByName = this.fields.stream()
				.collect(Collectors.toMap(f -> f.getFieldSpec().getName(), Function.identity()));
	}

	public List<Segment> getSegments() {
		return segments;
	}

	public List<FieldDefinition> getFields() {
		return fields;
	}

	public FieldDefinition getFieldByName(String name) {
		return fieldsByName.get(name);
	}

	public static class FieldDefinition {
		private final int column;
		private final FieldSpec fieldSpec;

		public FieldDefinition(int column, FieldSpec fieldSpec) {
			this.column = column;
			this.fieldSpec = fieldSpec;
		}

		public int getColumn() {
			return column;
		}

		public FieldSpec getFieldSpec() {
			return fieldSpec;
		}
	}

	public static DataSet makeRandomDataSet(int rows, int segmentSize, FieldSpec... fields) {
		final Random rng = new Random();
		final List<Segment> segments = new ArrayList<>();
		final int nSegments = rows / segmentSize;
		for (int i = 0; i < nSegments; i++) {
			segments.add(makeRandomSegment(rng, fields, segmentSize));
		}
		segments.add(makeRandomSegment(rng, fields, rows % segmentSize));
		return new DataSet(Arrays.asList(fields), segments);
	}

	private static Segment makeRandomSegment(Random rng, FieldSpec[] fields, int segmentSize) {
		Segment segment = new Segment(fields.length);
		for (int col = 0; col < fields.length; col++) {
			final FieldSpec field = fields[col];
			segment.rawData[col] = rng.ints(segmentSize, field.getOrigin(), field.getBound()).toArray();
		}
		return segment;
	}
}


<|start_filename|>src/main/java/com/medallia/codegen/DiscardableClassLoader.java<|end_filename|>
package com.medallia.codegen;

/** Classloader that can be garbage collected after it's only class is unloaded */
public class DiscardableClassLoader {
	/**
	 * This method loads a single class in its own class loader.
	 * The classloader will be a child of the baseClass' classloader.
	 */
	public static <T> Class<? extends T> classFromBytes(final Class<T> baseClass, final String name, final byte[] bytecode) {
		return new ClassLoader(baseClass.getClassLoader()) {
			Class<? extends T> c = defineClass(name, bytecode, 0, bytecode.length).asSubclass(baseClass);
		}.c;
	}
}


<|start_filename|>src/main/java/com/medallia/dsl/ast/AndExpr.java<|end_filename|>
package com.medallia.dsl.ast;

public class AndExpr extends BinaryExpr {
	public AndExpr(Expr left, Expr right) {
		super(left, right);
	}

	@Override
	public <T> T visit(ExprVisitor<T> visitor) { return visitor.visit(this); }
}


<|start_filename|>src/test/java/com/medallia/codegen/JavaCodeGeneratorTest.java<|end_filename|>
package com.medallia.codegen;


import org.junit.Test;

import java.util.function.Supplier;

import static org.hamcrest.CoreMatchers.is;
import static org.junit.Assert.assertThat;

public class JavaCodeGeneratorTest {
	@Test
	public void testCompile() throws IllegalAccessException, InstantiationException {
		Class<? extends Supplier> compiled = SimpleJavaCompiler.compile(Supplier.class, fileCg -> {
			fileCg.generateImport(Supplier.class);

			fileCg.publicClass("Test").implement("Supplier<String>")
					.build(classCg -> {
						classCg.publicMethod("String", "get")
								.build(cg -> cg.print("return ").printQuoted("Hello").println(";"));
					});
		});
		Supplier supplier = compiled.newInstance();
		assertThat(supplier.get(), is("Hello"));
	}
}

<|start_filename|>src/main/java/com/medallia/dsl/compiler/QueryCompiler.java<|end_filename|>
package com.medallia.dsl.compiler;

import com.medallia.codegen.JavaCodeGenerator;
import com.medallia.codegen.SimpleJavaCompiler;
import com.medallia.data.DataSet;
import com.medallia.data.DataSet.FieldDefinition;
import com.medallia.data.Segment;
import com.medallia.dsl.FieldStats;
import com.medallia.dsl.Query;
import com.medallia.dsl.ast.Agg;
import com.medallia.dsl.ast.AggVisitor;
import com.medallia.dsl.ast.AndExpr;
import com.medallia.dsl.ast.ConstantExpr;
import com.medallia.dsl.ast.ExprVisitor;
import com.medallia.dsl.ast.InExpr;
import com.medallia.dsl.ast.NotExpr;
import com.medallia.dsl.ast.OrExpr;
import com.medallia.dsl.ast.StatsAggregate;

import java.lang.reflect.Constructor;
import java.lang.reflect.InvocationTargetException;
import java.util.function.Consumer;
import java.util.function.Supplier;
import java.util.stream.IntStream;

public class QueryCompiler<T> {
	protected final Query<T> query;
	protected final DataSet dataSet;

	public QueryCompiler(Query<T> query, DataSet dataSet) {
		this.query = query;
		this.dataSet = dataSet;
	}


	@SuppressWarnings("unchecked")
	public Supplier<CompiledQueryBase<T>>  compile() {
		final Agg<T> aggregate = query.aggregateOp.buildAgg();

		Class<? extends CompiledQueryBase<T>> compiled = (Class<? extends CompiledQueryBase<T>>)
				SimpleJavaCompiler.compile(CompiledQueryBase.class, fileCg -> {
			fileCg.generateImport(CompiledQueryBase.class);
			fileCg.generateImport(Segment.class);
			fileCg.generateImport(FieldStats.class);

			fileCg.publicClass("CompiledQuery")
					.extend("CompiledQueryBase")
					.build(classCg -> {
						classCg.declareExisting("result");

						classCg.println("public CompiledQuery(Object result)");
						classCg.beginBlock();
						classCg.println("this.result = result;");
						classCg.endBlock();

						classCg.publicMethod("void", "process").arg("Segment", "segment")
								.build(methodCg -> {
									methodCg.declare("int[][]", "rawData", "segment.rawData");
									methodCg.declare("int", "nRows", "rawData[0].length");
									final String resultType = resultType(aggregate);
									methodCg.declare(resultType, "result", "(%s)this.result", resultType);
									methodCg.boundedLoop("row", "0", "nRows", loopCg -> {
										loopCg.ifThen(this::generateFilter, this::generateAggregate);
									});
								});
					});

		});
		Constructor<? extends CompiledQueryBase<T>> constructor;
		try {
			constructor = compiled.getConstructor(Object.class);
		} catch (NoSuchMethodException e) {
			throw new Error("Cannot find constructor", e);
		}
		return () -> {
			try {
				return constructor.newInstance(aggregate.makeResult(dataSet));
			} catch (InstantiationException | IllegalAccessException | InvocationTargetException e) {
				throw new RuntimeException("unexpected exception instantiating query", e);
			}
		};
	}

	protected void generateFilter(JavaCodeGenerator cg) {
		cg.print(query.buildExpressionTree().visit(new ExprVisitor<String>() {
			@Override
			public String visit(AndExpr andExpr) {
				return "(" + andExpr.getLeft().visit(this) + " && " + andExpr.getRight().visit(this) + ")";
			}

			@Override
			public String visit(InExpr inExpr) {
				return generateInExpr(cg, inExpr);
			}

			@Override
			public String visit(NotExpr notExpr) {
				return "!(" + notExpr.getTarget().visit(this) + ")";
			}

			@Override
			public String visit(OrExpr orExpr) {
				return "(" + orExpr.getLeft().visit(this) + " || " + orExpr.getRight().visit(this) + ")";
			}

			@Override
			public String visit(ConstantExpr constantExpr) {
				return String.valueOf(constantExpr.value);
			}
		}));
	}

	protected String generateInExpr(JavaCodeGenerator cg, InExpr inExpr) {
		FieldDefinition field = dataSet.getFieldByName(inExpr.getFieldName());
		return "(" + IntStream.of(inExpr.getValues())
				.mapToObj(v -> String.format("rawData[%d][row] == %dL", field.getColumn(), v))
				.reduce((a,b) -> a + " || " + b)
				.orElseThrow(() -> new RuntimeException("empty filter"))
				+ ")";
	}

	private void generateAggregate(JavaCodeGenerator cg) {
		Consumer<JavaCodeGenerator> generator = query.aggregateOp.buildAgg().visit(new AggVisitor<Consumer<JavaCodeGenerator>>() {
			@Override
			public Consumer<JavaCodeGenerator> visit(StatsAggregate statsAggregate) {
				return cg -> {
					cg.println(cg.variable("result") + ".count++;");
					cg.printf("%s.sum += rawData[%d][row];%n", cg.variable("result"), dataSet.getFieldByName(statsAggregate.getFieldName()).getColumn());
				};
			}

		});

		generator.accept(cg);
	}

	private String resultType(Agg<?> agg) {
		return agg.visit(statsAggregate -> "FieldStats");
	}
}


<|start_filename|>src/main/java/com/medallia/dsl/Query.java<|end_filename|>
package com.medallia.dsl;

import com.medallia.dsl.ast.Expr;

import java.util.List;

public class Query<R> {
	public final List<ConditionalExpression> filters;
	public final Aggregate<R> aggregateOp;

	public Query(List<ConditionalExpression> filters, Aggregate<R> aggregateOp) {
		this.filters = filters;
		this.aggregateOp = aggregateOp;
	}

	public Expr buildExpressionTree() {
		return filters.stream()
				.reduce(ConditionalExpression::and)
				.orElseGet(() -> ConditionalExpression.constant(true))
				.buildTree();
	}
}


<|start_filename|>src/main/java/com/medallia/dsl/ast/OrExpr.java<|end_filename|>
package com.medallia.dsl.ast;

public class OrExpr extends BinaryExpr {
	public OrExpr(Expr left, Expr right) {
		super(left, right);
	}

	@Override
	public <T> T visit(ExprVisitor<T> visitor) { return visitor.visit(this); }
}


<|start_filename|>src/main/java/com/medallia/benchmark/InterpretedQueryBenchmark.java<|end_filename|>
package com.medallia.benchmark;

import com.medallia.dsl.FieldStats;
import com.medallia.dsl.interpreter.QueryInterpreter;

public class InterpretedQueryBenchmark
		extends QueryBenchmark<QueryInterpreter<FieldStats>>
{
	public InterpretedQueryBenchmark() {
		super(((dataSet, query) -> new QueryInterpreter<>(query)),
				(dataSet, interpreter) -> interpreter.eval(dataSet));
	}
}


<|start_filename|>src/main/java/com/medallia/dsl/ast/BinaryExpr.java<|end_filename|>
package com.medallia.dsl.ast;

public abstract class BinaryExpr implements Expr {
	protected final Expr left;
	protected final Expr right;

	public BinaryExpr(Expr left, Expr right) {

		this.left = left;
		this.right = right;
	}

	public Expr getLeft() {
		return left;
	}

	public Expr getRight() {
		return right;
	}
}


<|start_filename|>src/main/java/com/medallia/benchmark/CompiledQueryBenchmark.java<|end_filename|>
package com.medallia.benchmark;

import com.medallia.dsl.FieldStats;
import com.medallia.dsl.compiler.CompiledQueryBase;
import com.medallia.dsl.compiler.QueryCompiler;

import java.util.function.Supplier;

public class CompiledQueryBenchmark
		extends QueryBenchmark<Supplier<CompiledQueryBase<FieldStats>>>
{
	public CompiledQueryBenchmark() {
		super(((dataSet, query) -> new QueryCompiler<>(query, dataSet).compile()), (dataSet, supplier) -> {
			CompiledQueryBase<FieldStats> compiledQuery = supplier.get();
			dataSet.getSegments().forEach(compiledQuery::process);
			return compiledQuery.getResult();
		});
	}
}


<|start_filename|>src/main/java/com/medallia/dsl/interpreter/QueryInterpreter.java<|end_filename|>
package com.medallia.dsl.interpreter;

import com.medallia.data.DataSet;
import com.medallia.data.Segment;
import com.medallia.dsl.FieldStats;
import com.medallia.dsl.Query;
import com.medallia.dsl.ast.Agg;
import com.medallia.dsl.ast.AggVisitor;
import com.medallia.dsl.ast.AndExpr;
import com.medallia.dsl.ast.ConstantExpr;
import com.medallia.dsl.ast.Expr;
import com.medallia.dsl.ast.ExprVisitor;
import com.medallia.dsl.ast.InExpr;
import com.medallia.dsl.ast.NotExpr;
import com.medallia.dsl.ast.OrExpr;
import com.medallia.dsl.ast.StatsAggregate;

public class QueryInterpreter<T> {

	private final Query<T> query;

	public QueryInterpreter(Query<T> query) {
		this.query = query;
	}

	public T eval(DataSet dataSet) {
		final Expr expr = query.buildExpressionTree();
		final Agg<T> agg = query.aggregateOp.buildAgg();
		final T result = agg.makeResult(dataSet);
		final Aggregator aggregator = makeAggregator(dataSet, agg);

		for (Segment segment : dataSet.getSegments()) {
			final int[][] rawData = segment.rawData;
			final int nRows = rawData[0].length;
			for (int row = 0; row < nRows; row++) {
				if (eval(expr, dataSet, segment, row)) {
					aggregator.process(segment, row, result);
				}
			}
		}

		return result;
	}

	private boolean eval(Expr expression, DataSet dataSet, Segment segment, int row) {
		return expression.visit(new ExprVisitor<Boolean>() {
			@Override
			public Boolean visit(AndExpr andExpr) {
				return andExpr.getLeft().visit(this) && andExpr.getRight().visit(this);
			}

			@Override
			public Boolean visit(InExpr inExpr) {
				int column = dataSet.getFieldByName(inExpr.getFieldName()).getColumn();
				int[] values = inExpr.getValues();
				for (int value : values) {
					if (value == segment.rawData[column][row])
						return true;
				}
				return false;
			}

			@Override
			public Boolean visit(NotExpr notExpr) {
				return !notExpr.getTarget().visit(this);
			}

			@Override
			public Boolean visit(OrExpr orExpr) {
				return orExpr.getLeft().visit(this) || orExpr.getRight().visit(this);
			}

			@Override
			public Boolean visit(ConstantExpr constantExpr) {
				return constantExpr.value;
			}
		});
	}

	static Aggregator makeAggregator(DataSet dataSet, Agg<?> agg) {
		return agg.visit(new AggVisitor<Aggregator>() {
			@Override
			public Aggregator visit(StatsAggregate statsAggregate) {
				int column = dataSet.getFieldByName(statsAggregate.getFieldName())
						.getColumn();
				return (segment, row, result) -> {
					((FieldStats)result).count++;
					((FieldStats)result).sum += segment.rawData[column][row];
				};
			}
		});
	}

	interface Aggregator {
		void process(Segment segment, int row, Object result);
	}
}


<|start_filename|>src/main/java/com/medallia/dsl/FieldStats.java<|end_filename|>
package com.medallia.dsl;

import java.util.Objects;

public final class FieldStats {
	public long count;
	public double sum;

	@Override
	public String toString() {
		return String.format("Stats[count: %d, sum: %s, avg: %s]", count, sum, sum/count);
	}

	@Override
	public int hashCode() {
		return Objects.hash(count, sum);
	}

	@Override
	public boolean equals(Object obj) {
		return obj != null
				&& obj.getClass() == getClass()
				&& ((FieldStats)obj).count == count
				&& ((FieldStats)obj).sum   == sum;
	}
}


<|start_filename|>src/main/java/com/medallia/dsl/ast/StatsAggregate.java<|end_filename|>
package com.medallia.dsl.ast;

import com.medallia.data.DataSet;
import com.medallia.data.Segment;
import com.medallia.dsl.Aggregate;
import com.medallia.dsl.FieldStats;

public class StatsAggregate implements Agg<FieldStats> {
	private final String fieldName;
	public StatsAggregate(String fieldName) {
		this.fieldName = fieldName;
	}

	public String getFieldName() {
		return fieldName;
	}

	@Override
	public FieldStats makeResult(DataSet dataSet) {
		return new FieldStats();
	}

	@Override
	public <T> T visit(AggVisitor<T> visitor) {
		return visitor.visit(this);
	}
}


<|start_filename|>src/main/java/com/medallia/benchmark/QueryBenchmark.java<|end_filename|>
package com.medallia.benchmark;

import com.medallia.codegen.SimpleJavaCompiler;
import com.medallia.data.DataSet;
import com.medallia.data.FieldSpec;
import com.medallia.dsl.FieldStats;
import com.medallia.dsl.Query;
import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.Level;
import org.openjdk.jmh.annotations.Measurement;
import org.openjdk.jmh.annotations.OutputTimeUnit;
import org.openjdk.jmh.annotations.Scope;
import org.openjdk.jmh.annotations.Setup;
import org.openjdk.jmh.annotations.State;
import org.openjdk.jmh.annotations.Warmup;
import org.openjdk.jmh.infra.Blackhole;
import org.openjdk.jmh.runner.Runner;
import org.openjdk.jmh.runner.RunnerException;
import org.openjdk.jmh.runner.options.Options;
import org.openjdk.jmh.runner.options.OptionsBuilder;

import java.util.concurrent.TimeUnit;

import static com.medallia.dsl.Aggregate.statsAggregate;
import static com.medallia.dsl.ConditionalExpression.field;
import static com.medallia.dsl.ConditionalExpression.not;
import static com.medallia.dsl.QueryBuilder.newQuery;

/**
 * Base query benchmark.
 * Query execution strategy is provided by two functional interfaces.
 * It provides two benchmarks, one for simple and one for complex queries (CPU intensive).
 * This benchmark is single-threaded.
 * @param <T> type of the state (e.g. the compiled query for compiling runners)
 */
@State(Scope.Benchmark)
@Warmup(iterations = 1)
@Measurement(iterations = 3)
@OutputTimeUnit(TimeUnit.SECONDS)
public abstract class QueryBenchmark<T> {

	private final StateInitializer<T> stateInitializer;
	private final QueryRunner<T> queryRunner;
	private final DataSet dataSet;

	private T state;

	protected QueryBenchmark(StateInitializer<T> stateInitializer, QueryRunner<T> queryRunner) {
		this.stateInitializer = stateInitializer;
		this.queryRunner = queryRunner;
		dataSet = DataSet.makeRandomDataSet(
				1_000_000, // rows
				50_000,	   // segment size
				new FieldSpec("a", 0, 11),
				new FieldSpec("b", 0, 5),
				new FieldSpec("c", 0, 11),
				new FieldSpec("d", 0, 5),
				new FieldSpec("e", 0, 20),
				new FieldSpec("f", 0, 5),
				new FieldSpec("ltr", 0, 11)
		);
	}

	@Setup(Level.Trial)
	public void setup() {
		state = null;
	}

	private void initializeState(Query<FieldStats> query) {
		if (state == null) {
			// this will happen only on the warmup iteration
			state = stateInitializer.initialize(dataSet, query);
		}
	}

	@Benchmark
	public void simpleQuery(Blackhole blackhole) {
		final Query<FieldStats> query = newQuery()
				.filter(field("a").in(1, 2, 3).or(field("b").is(3)))
				.aggregate(statsAggregate("ltr"));

		initializeState(query);

		blackhole.consume(queryRunner.run(dataSet, state));
	}

	@Benchmark
	public void complexQuery(Blackhole blackhole) {
		final Query<FieldStats> query = newQuery()
				.filter(field("a").in(1, 2, 3, 5).or(field("b").is(3)))
				.filter(field("c").in(3, 5, 8, 1, 2).or(field("d").in(1,3,4)))
				.filter(not(field("e").in(2, 5, 8,10,12,13,14,15,17)).or(field("f").in(1,4,2)))
				.aggregate(statsAggregate("ltr"));

		initializeState(query);

		blackhole.consume(queryRunner.run(dataSet, state));
	}

	@FunctionalInterface
	public interface StateInitializer<T> {
		T initialize(DataSet dataSet, Query<FieldStats> query);
	}

	@FunctionalInterface
	public interface QueryRunner<T> {
		FieldStats run(DataSet dataSet, T state);
	}

	public static void main(String[] args) throws RunnerException {
		// Globally enable dumping the source code for simplicity
		SimpleJavaCompiler.dump=true;

		Options opts = new OptionsBuilder()
				.include(CompiledQueryBenchmark.class.getSimpleName())
				.include(VariantCompiledQueryBenchmark.class.getSimpleName())
				.include(InterpretedQueryBenchmark.class.getSimpleName())
				.include(StreamInterpretedQueryBenchmark.class.getSimpleName())
				.forks(2)
				.build();
		new Runner(opts).run();
	}
}


<|start_filename|>src/main/java/com/medallia/benchmark/UnsafeCompiledQueryBenchmark.java<|end_filename|>
package com.medallia.benchmark;

import com.medallia.dsl.FieldStats;
import com.medallia.dsl.compiler.BranchReducingQueryCompiler;
import com.medallia.dsl.compiler.CompiledQueryBase;
import com.medallia.dsl.compiler.unsafe.CCompiledQuery;
import com.medallia.dsl.compiler.unsafe.CQueryCompiler;
import org.openjdk.jmh.runner.Runner;
import org.openjdk.jmh.runner.RunnerException;
import org.openjdk.jmh.runner.options.Options;
import org.openjdk.jmh.runner.options.OptionsBuilder;

import java.util.function.Supplier;

public class UnsafeCompiledQueryBenchmark
		extends QueryBenchmark<Supplier<CompiledQueryBase<FieldStats>>>
{
	public UnsafeCompiledQueryBenchmark() {
		super(((dataSet, query) -> new CQueryCompiler<>(query, dataSet).compile()), (dataSet, supplier) -> {
			CompiledQueryBase<FieldStats> compiledQuery = supplier.get();
			dataSet.getSegments().forEach(compiledQuery::process);
			return compiledQuery.getResult();
		});
	}

	public static void main(String[] args) throws RunnerException {
		Options opts = new OptionsBuilder()
				.include(UnsafeCompiledQueryBenchmark.class.getSimpleName())
				.forks(2)
				.build();
		new Runner(opts).run();
	}
}


<|start_filename|>src/main/java/com/medallia/dsl/ast/UnaryExpr.java<|end_filename|>
package com.medallia.dsl.ast;

public abstract class UnaryExpr implements Expr {
	protected final Expr target;

	public UnaryExpr(Expr target) {
		this.target = target;
	}

	public Expr getTarget() {
		return target;
	}
}


<|start_filename|>src/main/java/com/medallia/dsl/ast/ExprVisitor.java<|end_filename|>
package com.medallia.dsl.ast;

/** Visitor interface for expressions */
public interface ExprVisitor<T> {
	T visit(AndExpr andExpr);

	T visit(InExpr inExpr);

	T visit(NotExpr notExpr);

	T visit(OrExpr orExpr);

	T visit(ConstantExpr constantExpr);
}


<|start_filename|>src/main/java/com/medallia/dsl/ast/NotExpr.java<|end_filename|>
package com.medallia.dsl.ast;

public class NotExpr extends UnaryExpr {
	public NotExpr(Expr target) {
		super(target);
	}

	@Override
	public <T> T visit(ExprVisitor<T> visitor) { return visitor.visit(this); }
}


<|start_filename|>src/main/java/com/medallia/dsl/ConditionalExpression.java<|end_filename|>
package com.medallia.dsl;

import com.medallia.dsl.ast.AndExpr;
import com.medallia.dsl.ast.ConstantExpr;
import com.medallia.dsl.ast.Expr;
import com.medallia.dsl.ast.InExpr;
import com.medallia.dsl.ast.NotExpr;
import com.medallia.dsl.ast.OrExpr;

import java.util.function.Supplier;

public class ConditionalExpression {
	private final Supplier<Expr> nodeBuilder;

	private ConditionalExpression(Supplier<Expr> nodeBuilder) {
		this.nodeBuilder = nodeBuilder;
	}

	public Expr buildTree() { return nodeBuilder.get(); }

	public ConditionalExpression and(ConditionalExpression other) {
		return new ConditionalExpression(() -> new AndExpr(nodeBuilder.get(), other.nodeBuilder.get()));
	}

	public ConditionalExpression or(ConditionalExpression other) {
		return new ConditionalExpression(() -> new OrExpr(nodeBuilder.get(), other.nodeBuilder.get()));
	}

	public static ConditionalExpression not(ConditionalExpression other) {
		return new ConditionalExpression(() -> new NotExpr(other.nodeBuilder.get()));
	}

	public static FieldExpressionBuilder field(String fieldName) {
		return new FieldExpressionBuilder(fieldName);
	}

	public static ConditionalExpression constant(boolean value) {
		return new ConditionalExpression(() -> new ConstantExpr(value));
	}

	public static class FieldExpressionBuilder {
		private final String fieldName;

		FieldExpressionBuilder(String fieldName) {
			this.fieldName = fieldName;
		}

		public final ConditionalExpression is(int value) {
			return in(value);
		}

		public final ConditionalExpression in(int... values) {
			return new ConditionalExpression(() -> new InExpr(fieldName, values));
		}
	}
}

<|start_filename|>src/main/java/com/medallia/benchmark/StreamInterpretedQueryBenchmark.java<|end_filename|>
package com.medallia.benchmark;

import com.medallia.dsl.FieldStats;
import com.medallia.dsl.interpreter.StreamQueryInterpreter;

public class StreamInterpretedQueryBenchmark
		extends QueryBenchmark<StreamQueryInterpreter<FieldStats>>
{
	public StreamInterpretedQueryBenchmark() {
		super(((dataSet, query) -> new StreamQueryInterpreter<>(query)),
				(dataSet, interpreter) -> interpreter.eval(dataSet));
	}
}


<|start_filename|>src/main/java/com/medallia/dsl/ast/InExpr.java<|end_filename|>
package com.medallia.dsl.ast;

public class InExpr implements Expr {
	private final String fieldName;
	private final int[] values;

	public <T> InExpr(String fieldName, int[] values) {
		this.fieldName = fieldName;
		this.values = values;
	}

	public String getFieldName() {
		return fieldName;
	}

	public int[] getValues() {
		return values;
	}

	@Override
	public <T> T visit(ExprVisitor<T> visitor) { return visitor.visit(this); }

}


<|start_filename|>src/main/java/com/medallia/dsl/ast/ConstantExpr.java<|end_filename|>
package com.medallia.dsl.ast;

public class ConstantExpr implements Expr {
	public final boolean value;

	public ConstantExpr(boolean value) {
		this.value = value;
	}

	@Override
	public <T> T visit(ExprVisitor<T> visitor) {
		return visitor.visit(this);
	}
}


<|start_filename|>src/main/java/com/medallia/Main.java<|end_filename|>
package com.medallia;

import com.medallia.data.DataSet;
import com.medallia.data.FieldSpec;
import com.medallia.dsl.FieldStats;
import com.medallia.dsl.Query;
import com.medallia.dsl.compiler.BranchReducingQueryCompiler;
import com.medallia.dsl.compiler.CompiledQueryBase;
import com.medallia.dsl.compiler.QueryCompiler;
import com.medallia.dsl.interpreter.QueryInterpreter;
import com.medallia.dsl.interpreter.StreamQueryInterpreter;

import static com.medallia.dsl.Aggregate.statsAggregate;
import static com.medallia.dsl.ConditionalExpression.field;
import static com.medallia.dsl.QueryBuilder.newQuery;

public class Main {
	public static void main(String[] args) {
		DataSet dataSet = DataSet.makeRandomDataSet(
				1_000_000, // rows
				50_000,	   // segment size
				new FieldSpec("a", 0, 11),
				new FieldSpec("b", 0, 5),
				new FieldSpec("ltr", 0, 11)
				);

		distribution(dataSet);
	}

	private static void distribution(DataSet dataSet) {
		Query<FieldStats> query = newQuery()
				.filter(
						field("a").in(1, 2, 3)
								.or(field("b").is(3))
				)
				.aggregate(statsAggregate("ltr"));


		final QueryInterpreter<FieldStats> interpreter = new QueryInterpreter<>(query);
		FieldStats interpreterResult = interpreter.eval(dataSet);

		final QueryCompiler<FieldStats> queryCompiler = new QueryCompiler<>(query, dataSet);
		final CompiledQueryBase<FieldStats> compiledQuery = queryCompiler.compile().get();
		dataSet.getSegments().forEach(compiledQuery::process);
		FieldStats compiledResult = compiledQuery.getResult();

		final StreamQueryInterpreter<FieldStats> streamQueryInterpreter = new StreamQueryInterpreter<>(query);
		FieldStats streamResult = streamQueryInterpreter.eval(dataSet);

		final BranchReducingQueryCompiler<FieldStats> branchReducingQueryCompiler = new BranchReducingQueryCompiler<>(query, dataSet);
		final CompiledQueryBase<FieldStats> compiledQuery2 = branchReducingQueryCompiler.compile().get();
		dataSet.getSegments().forEach(compiledQuery2::process);
		FieldStats branchResult = compiledQuery2.getResult();

		System.out.println("            Interpreted result: " + interpreterResult);
		System.out.println("               Compiled result: " + compiledResult);
		System.out.println("     Stream Interpreted result: " + streamResult);
		System.out.println("Branch reduced compiled result: " + branchResult);
	}
}


<|start_filename|>src/main/java/com/medallia/codegen/JavaCodeGenerator.java<|end_filename|>
package com.medallia.codegen;

import java.io.PrintWriter;
import java.util.Arrays;
import java.util.StringJoiner;
import java.util.function.Consumer;
import com.medallia.io.IndentedPrintWriter;

import java.io.StringWriter;
import java.io.Writer;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Stream;


/**
 * Utility class to generate Java source code.
 * Each instance of {@link JavaCodeGenerator} defines a scope. This is used to keep track
 * of names declared in each scope and allows for variable name shadowing. A new scope is
 * created on blocks and for-loops following java scoping rules.
 * Scopes also collect code emitted in them, once the scope is left, all code emitted is
 * applied to its parent scope.
 * The {@link JavaCodeGenerator} also define a parent-child relationship, where each child
 * defines a new scope based on the current scope of it's parent.
 * This is used to support, for example, defining new methods/fields at a class level while
 * generating code for a method.
 * Another way to see it is that considering the scopes defined by a java compilation unit,
 * each JavaCodeGenerator can generate code on a specific scope subtree. So, for example,
 * one can have a code generator for the class, and one for each method.
 */
@SuppressWarnings({"WeakerAccess", "unused"})
public class JavaCodeGenerator {
	/** Current scope */
	private Scope scope;

	/** Parent of this code generator */
	private final JavaCodeGenerator parent;

	/** Creates a new JavaCodeGenerator that writes to the specified writer */
	public JavaCodeGenerator() {
		scope = new Scope();
		parent = null;
	}

	/** Creates a new JavaCodeGenerator that generates in a sub-scope of the specified generator */
	private JavaCodeGenerator(JavaCodeGenerator parent) {
		this.scope = new Scope(parent.scope);
		this.parent = parent;
	}

	/** @return the name with a unique integer appended. The number is unique only for the current scope. */
	public String unique(String name) {
		return scope.unique(name);
	}

	/** @return this code generator's parent or null if top-level */
	public JavaCodeGenerator parent() {
		return parent;
	}

	/** Generates an import */
	public void generateImport(Class<?> aClass) {
		printf("import %s;%n", aClass.getName());
	}

	/** Declares an uninitialized local variable */
	public JavaCodeGenerator declare(String type, String name) {
		return declare(type, name, null);
	}

	/** Declare a pre-existing variable in the current scope. No code is emitted. */
	public JavaCodeGenerator declareExisting(String name) {
		scope.declare(name);
		return this;
	}

	/** Declares an initialized local variable */
	public JavaCodeGenerator declare(String type, String name, String initialValue, Object... args) {
		scope.delayed.add(declareInternal(scope, type, name, initialValue == null ? null : String.format(initialValue, args)));
		return this;
	}

	/**
	 * Declares an initialized local variable at the beginning of the current scope.
	 * Useful to declare temporary variables on complex expressions.
	 */
	public JavaCodeGenerator declareOnScope(final String type, final String name, final String initialValue) {
		scope.declare(type, name, initialValue);
		return this;
	}

	/**
	 * Declares an initialized local variable at the beginning of the current scope.
	 * Useful to declare temporary variables on complex expressions.
	 * @return the resolved name of the variable
	 */
	public static String declareOnScope(final Scope scope, final String type, final String name, final String initialValue) {
		scope.delayed.add(scope.lastDeclaration++, declareInternal(scope, type, name, initialValue));
		return scope.variable(name);
	}

	/**
	 * Prints the specified code at the beginning of the current scope.
	 * Useful to initialize temporary variables on complex expressions.
	 */
	public JavaCodeGenerator printfOnScope(final String format, final Object... args) {
		scope.delayed.add(scope.lastDeclaration++, printfInternal(format, args));
		return this;
	}

	private Consumer<IndentedPrintWriter> printfInternal(final String format, final Object[] args) {
		return pw -> pw.printf(format, args);
	}

	private static Consumer<IndentedPrintWriter> declareInternal(final Scope scope, final String type, String name, final String initialValue) {
		scope.declare(name);
		final String resolvedVariable = scope.variable(name);
		return pw -> {
			if (initialValue == null) {
				pw.printf("%s %s;%n", type, resolvedVariable);
			} else {
				pw.printf("%s %s = %s;%n", type, resolvedVariable, initialValue);
			}
		};
	}

	/** Simple bounded loop over an int index variable */
	public JavaCodeGenerator boundedLoop(String indexVar, String lowerBound, String upperBound, Consumer<JavaCodeGenerator> body) {
		enterScope();
		scope.declare(indexVar);
		printf("for (int %1$s = %2$s; %1$s < %3$s; %1$s++) {%n", variable(indexVar), lowerBound, upperBound);
		indent();
		body.accept(this);
		dedent();
		println("}");
		leaveScope();
		return this;
	}

	/** Unbounded loop */
	public JavaCodeGenerator loop(Consumer<JavaCodeGenerator> body) {
		println("for (;;) {");
		enterScope();
		indent();
		body.accept(this);
		dedent();
		println("}");
		leaveScope();
		return this;
	}

	/** Fore-each loop over an iterable */
	public JavaCodeGenerator forEach(String type, String elementVar, String collection, Consumer<JavaCodeGenerator> body) {
		enterScope();
		scope.declare(elementVar);
		printf("for (%1$s %2$s : %3$s) {%n", type, variable(elementVar), collection);
		indent();
		body.accept(this);
		dedent();
		println("}");
		leaveScope();
		return this;
	}

	/** Emmit a single-line comment. Supports new lines. */
	public void comment(String s, Object... args) {
		println("//" + String.format(s, args).replaceAll("\n", "\n//"));
	}

	/** @return the mangled name of the variable if it's shadowing another on an outer scope */
	public String variable(String name) {
		return scope.variable(name);
	}

	/** Print the specified string so it is a valid java string literal. */
	public JavaCodeGenerator printQuoted(String s) {
		if (s == null) {
			print("null");
		} else {
			print('"');
			print(escape(s));
			print('"');
		}
		return this;
	}

	/** Generates a simple if-then statement */
	public JavaCodeGenerator ifThen(final String condition, Consumer<JavaCodeGenerator> accept) {
		return ifThen(literal(condition), accept);
	}

	/** Generates a simple if-then statement */
	public JavaCodeGenerator ifThen(Consumer<JavaCodeGenerator> condition, Consumer<JavaCodeGenerator> accept) {
		return ifElse(condition, accept, null);
	}

	/** Generate an if-else statement */
	public JavaCodeGenerator ifElse(String condition, Consumer<JavaCodeGenerator> accept, Consumer<JavaCodeGenerator> reject) {
		return ifElse(literal(condition), accept, reject);
	}

	/** @return the current scope */
	public Scope getScope() {
		return scope;
	}

	/** Generate an if-else statement */
	public JavaCodeGenerator ifElse(Consumer<JavaCodeGenerator> condition, Consumer<JavaCodeGenerator> accept, Consumer<JavaCodeGenerator> reject) {
		print("if (");
		condition.accept(this);
		println(") {");
		enterScope();
		indent();
		accept.accept(this);
		dedent();
		leaveScope();
		if (reject != null) {
			println("} else {");
			enterScope();
			indent();
			reject.accept(this);
			dedent();
			leaveScope();
		}
		println("}");
		return this;
	}

	private Consumer<JavaCodeGenerator> literal(final String condition) {
		return cg -> cg.print(condition);
	}

	/** Close all open scopes and generate the code */
	public void generate(Writer out) {
		final IndentedPrintWriter pw = new IndentedPrintWriter(out);
		leaveUntil(parent == null ? null : parent.scope);
		for (Consumer<IndentedPrintWriter> cmd : scope.delayed) {
			cmd.accept(pw);
		}
	}

	private void leaveUntil(Scope root) {
		while (scope.parent != root)
			leaveScope();
	}

	/** Begins a new block */
	public JavaCodeGenerator beginBlock() {
		println("{");
		indent();
		enterScope();
		return this;
	}

	/** Ends the current block */
	public JavaCodeGenerator endBlock() {
		leaveScope();
		dedent();
		println("}");
		return this;
	}

	private void enterScope() {
		scope = new Scope(scope);
	}

	private void leaveScope() {
		if (scope.parent == null) throw new IllegalStateException("you must be in a scope to leave it");
		scope.parent.delayed.addAll(scope.delayed);
		scope = scope.parent;
	}

	private String escape(String s) {
		final int len = s.length();
		StringBuilder sb = new StringBuilder(len);
		for (int i = 0; i < len; i++) {
			char c = s.charAt(i);
			switch (c) {
				case '"' : sb.append("\\\""); break;
				case '\\': sb.append("\\\\"); break;
				case '\n': sb.append("\\n") ; break;
				case '\r': sb.append("\\r") ; break;
				default  : sb.append(c)     ; break;
			}
		}
		return sb.toString();
	}

	/** Print a new line */
	public JavaCodeGenerator println() {
		scope.delayed.add(PrintWriter::println);
		return this;
	}

	/** Prints an Object and then terminates the line. */
	public JavaCodeGenerator println(final Object s) {
		scope.delayed.add(pw -> pw.println(s));
		return this;
	}

	/** Prints an object. */
	public JavaCodeGenerator print(final Object s) {
		scope.delayed.add(pw -> pw.print(s));
		return this;
	}

	/** Prints a formatted string using the specified format string and arguments. */
	public JavaCodeGenerator printf(final String format, final Object ... args) {
		scope.delayed.add(printfInternal(format, args));
		return this;
	}

	private void indent() {
		scope.delayed.add(IndentedPrintWriter::indent);
	}

	private void dedent() {
		scope.delayed.add(IndentedPrintWriter::dedent);
	}

	/** Build a new public class */
	public ClassBuilder publicClass(String name) {
		return new ClassBuilder("public", name);
	}

	/** Helper class to generate classes */
	public class ClassBuilder {
		private ClassBuilder(String visibility, String name) {
			declareExisting(name);
			if (visibility != null) {
				printf("%s class %s", visibility, name);
			} else {
				printf("class %s", name);
			}
		}

		/** Make the generated class extend another class. */
		public ClassBuilder extend(String aClass) {
			printf(" extends %s", aClass);
			return this;
		}

		/** Make the generated class extend another class. */
		public ClassBuilder implement(String... anInterface) {
			StringJoiner joiner = new StringJoiner(",");
			for (String s : anInterface) {
				joiner.add(s);
			}
			printf(" implements %s", joiner);
			return this;
		}

		/** Build the class' body */
		public JavaCodeGenerator build(Consumer<JavaCodeGenerator> body) {
			println(" {");
			enterScope();
			indent();
			body.accept(JavaCodeGenerator.this);
			dedent();
			println("}");
			leaveScope();
			return JavaCodeGenerator.this;
		}
	}

	/** Create a private method */
	public MethodBuilder packageMethod(String returnType, String name) {
		return new MethodBuilder("", returnType, name);
	}

	/** Create a private method */
	public MethodBuilder privateMethod(String returnType, String name) {
		return new MethodBuilder("private ", returnType, name);
	}

	/** Create a public method */
	public MethodBuilder publicMethod(String returnType, String name) {
		return new MethodBuilder("public ", returnType, name);
	}

	/** Helper class to build methods */
	public class MethodBuilder {
		private JavaCodeGenerator cg = new JavaCodeGenerator(JavaCodeGenerator.this);
		private boolean hasArgs;

		private MethodBuilder(String visibility, String returnType, String name) {
			cg.printf("%s%s %s(", visibility, returnType, name);
		}

		/** Adds an argument of the specified type */
		public MethodBuilder arg(String type, String name) {
			cg.declareExisting(name);
			if (hasArgs) {
				cg.print(", ");
			}
			cg.printf("%s %s", type, name);
			hasArgs = true;
			return this;
		}

		/** Builds this method's body */
		public JavaCodeGenerator build(Consumer<JavaCodeGenerator> body) {
			cg.println(") {");
			cg.indent();
			// Generate the body in it's own scope.
			final JavaCodeGenerator bodyGen = new JavaCodeGenerator(cg);
			body.accept(bodyGen);
			bodyGen.leaveScope();
			cg.dedent();
			cg.println("}");
			cg.leaveScope();
			return JavaCodeGenerator.this;
		}
	}

	/** Simple scope for generation of shadowed local variables */
	public static class Scope {
		private final Scope parent;
		private final Set<String> variables = new HashSet<>();
		private int lastDeclaration;
		private final List<Consumer<IndentedPrintWriter>> delayed = new ArrayList<>();
		private int varId;

		private Scope() {
			this.parent = null;
		}

		private Scope(Scope parent) {
			this.parent = parent;
		}

		private void declare(String name) {
			if(variables.contains(name)) throw new IllegalArgumentException(String.format("variable '%s' already declared in this scope", name));
			variables.add(name);
		}

		private String mangle(String name) {
			int level = -1;
			for (Scope scope = this; scope != null; scope = scope.parent) {
				if (scope.variables.contains(name)) {
					++level;
				}
			}
			return level > 0 ? name + "$" + level : name;
		}

		String variable(String name) {
			if (variables.contains(name)) {
				return mangle(name);
			} else if (parent != null) {
				return parent.variable(name);
			} else {
				throw new IllegalArgumentException("Unknown variable: " + name);
			}
		}

		private String unique(String name) {
			return name + "_" + (++varId);
		}

		/** Declares a variable at the beginning of the this scope. */
		public String declare(String type, String name, String initialValue) {
			delayed.add(lastDeclaration++, declareInternal(type, name, initialValue));
			return variable(name);
		}

		private Consumer<IndentedPrintWriter> declareInternal(final String type, String name, final String initialValue) {
			declare(name);
			final String resolvedVariable = variable(name);
			return pw -> {
				if (initialValue == null) {
					pw.printf("%s %s;%n", type, resolvedVariable);
				} else {
					pw.printf("%s %s = %s;%n", type, resolvedVariable, initialValue);
				}
			};
		}

		/** @return true if the specified variable is defined in this scope */
		public boolean isDefined(String variable) {
			return variables.contains(variable);
		}
	}
}

<|start_filename|>src/main/java/com/medallia/codegen/SimpleJavaCompiler.java<|end_filename|>
package com.medallia.codegen;

import javax.tools.FileObject;
import javax.tools.ForwardingJavaFileManager;
import javax.tools.JavaCompiler;
import javax.tools.JavaCompiler.CompilationTask;
import javax.tools.JavaFileObject;
import javax.tools.JavaFileObject.Kind;
import javax.tools.SimpleJavaFileObject;
import javax.tools.StandardJavaFileManager;
import javax.tools.ToolProvider;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.OutputStream;
import java.io.StringWriter;
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.function.Consumer;

/** Simple wrapper over javac to compile and class-load a single class */
public class SimpleJavaCompiler {

	private static final String CLASSPATH = System.getProperty("java.class.path");

	/** A javac instance */
	private final JavaCompiler compiler;
	/** The Standard file manager for the compiler */
	private final StandardJavaFileManager standardJavaFileManager;

	/** We keep an initialized copy for each thread. */
	private static final ThreadLocal<SimpleJavaCompiler> BY_THREAD = new ThreadLocal<SimpleJavaCompiler>() {
		@Override
		protected SimpleJavaCompiler initialValue() {
			return new SimpleJavaCompiler();
		}
	};

	/** Create a new compiler */
	private SimpleJavaCompiler() {
		compiler = ToolProvider.getSystemJavaCompiler();
		standardJavaFileManager = compiler.getStandardFileManager(null, null, null);
	}

	/**
	 * Compiles and class-loads a single class from the specified Java source code.
	 *
	 * To amortize the cost of initializing a JavaCompiler, underlying instances are cached
	 * on a per-thread basis. It is recommended that callers of this method reuse threads
	 * for increased performance and reduced memory consumption.
	 *
	 * @param source a String containing Java source code to generate the class
	 * @return the compiled and class-loaded class
	 */
	public static <T> Class<? extends T> compile(Class<T> superClass, String source) {
		return BY_THREAD.get().compile0(superClass, source);
	}


	/**
	 * Compiles and class-loads a single class from the specified Java source code.
	 *
	 * To amortize the cost of initializing a JavaCompiler, underlying instances are cached
	 * on a per-thread basis. It is recommended that callers of this method reuse threads
	 * for increased performance and reduced memory consumption.
	 *
	 * @param sourceGen a String containing Java source code to generate the class
	 * @return the compiled and class-loaded class
	 */
	public static <T> Class<? extends T> compile(Class<T> superClass, Consumer<JavaCodeGenerator> sourceGen) {
		final JavaCodeGenerator cg = new JavaCodeGenerator();
		sourceGen.accept(cg);

		final StringWriter sw = new StringWriter();
		cg.generate(sw);

		final Class<? extends T> compiledClass = compile(superClass, sw.toString());
		dump(sw, compiledClass);
		return compiledClass;
	}

	public static volatile boolean dump = Boolean.getBoolean("compiler.java.dump");

	public static volatile String dumpDir = System.getProperty("compiler.java.dump.dir");

	private static void dump(StringWriter sw, Class<?> compiledClass) {
		if (dump) {
			System.out.println(sw);
		}
		if (dumpDir != null) {
			try {
				Files.write(Paths.get(dumpDir, compiledClass.getCanonicalName().replace('.', File.separatorChar) + ".java"), sw.toString().getBytes());
			} catch (IOException e) {
				throw new RuntimeException("unable to dump source file", e);
			}
		}
	}


	/**
	 * Compiles and class-loads a single class from the specified Java source code.
	 * @param source a String containing Java source code to generate the class
	 * @return the compiled and class-loaded class
	 */
	private <T> Class<? extends T> compile0(Class<T> superClass, String source) {
		final List<JavaFileObject> compilationUnits = Collections.singletonList(new SourceFile(toUri("Generated"), source));
		final FileManager fileManager = new FileManager(standardJavaFileManager);
		final StringWriter output = new StringWriter();
		final CompilationTask task = compiler.getTask(output, fileManager, null, Arrays.asList("-g", "-proc:none", "-classpath", CLASSPATH), null, compilationUnits);

		if (!task.call()) {
			throw new IllegalArgumentException("Compilation failed:\n" + output + "\n Source code: \n" + source);
		}

		int classCount = fileManager.output.size();
		// 2nd condition to work around the case where additional "GuardedBy" generated in TA Spark tagging
		if (classCount == 1 || (classCount > 0 && "CompiledStage"
				.equals(fileManager.output.get(0).getName()))) {
			return DiscardableClassLoader.classFromBytes(superClass, null, fileManager.output.get(0).outputStream.toByteArray());
		}
		throw new IllegalArgumentException("Compilation yielded an unexpected number of classes: " + classCount);
	}

	/** Represents a source file whose contents is loaded from a String */
	private static class SourceFile extends SimpleJavaFileObject {
		private final String contents;

		private SourceFile(URI uri, String contents) {
			super(uri, Kind.SOURCE);
			this.contents = contents;
		}

		@Override
		public String getName() { return uri.getRawSchemeSpecificPart(); }

		/** Ignore the file name for public classes */
		@Override
		public boolean isNameCompatible(String simpleName, Kind kind) { return true; }

		@Override
		public CharSequence getCharContent(boolean ignoreEncodingErrors) throws IOException {
			return contents;
		}
	}

	/** A compiled class file. It's content is stored in a ByteArrayOutputStream. */
	private static class ClassFile extends SimpleJavaFileObject {
		private ByteArrayOutputStream outputStream;

		private ClassFile(URI uri) {
			super(uri, Kind.CLASS);
		}

		@Override
		public String getName() { return uri.getRawSchemeSpecificPart(); }

		@Override
		public OutputStream openOutputStream() throws IOException {
			return outputStream = new ByteArrayOutputStream();
		}
	}

	/** A simple file manager that collects written files in memory */
	private static class FileManager
			extends ForwardingJavaFileManager<StandardJavaFileManager>
	{

		private List<ClassFile> output = new ArrayList<>();

		FileManager(StandardJavaFileManager target) {
			super(target);
		}

		@Override
		public JavaFileObject getJavaFileForOutput(Location location, String className, Kind kind, FileObject sibling) {
			final ClassFile file = new ClassFile(toUri(className));
			output.add(file);
			return file;
		}
	}

	private static URI toUri(String path)  {
		try {
			return new URI(null, null, path, null);
		} catch (URISyntaxException e) {
			throw new RuntimeException("exception parsing uri", e);
		}
	}
}


<|start_filename|>src/main/java/com/medallia/data/Segment.java<|end_filename|>
package com.medallia.data;

public class Segment {
	public final int[][] rawData;

	public Segment(int cols) {
		rawData = new int[cols][];
	}
}


<|start_filename|>src/main/java/com/medallia/dsl/Aggregate.java<|end_filename|>
package com.medallia.dsl;

import com.medallia.dsl.ast.Agg;
import com.medallia.dsl.ast.StatsAggregate;

import java.util.function.Supplier;

public class Aggregate<R> {
	private final Supplier<Agg<R>> nodeBuilder;

	private Aggregate(Supplier<Agg<R>> nodeBuilder) {
		this.nodeBuilder = nodeBuilder;
	}

	public static Aggregate<FieldStats> statsAggregate(String fieldName) {
		return new Aggregate<>(() -> new StatsAggregate(fieldName));
	}

	public Agg<R> buildAgg() {
		return nodeBuilder.get();
	}
}

<|start_filename|>src/main/java/com/medallia/dsl/QueryBuilder.java<|end_filename|>
package com.medallia.dsl;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

public class QueryBuilder {

	private final List<ConditionalExpression> filters = new ArrayList<>();

	public static QueryBuilder newQuery() {
		return new QueryBuilder();
	}

	public QueryBuilder filter(ConditionalExpression expression) {
		filters.add(expression);
		return this;
	}

	public <R> Query<R> aggregate(Aggregate<R> aggregateOp) {
		return new Query<R>(Collections.unmodifiableList(filters), aggregateOp);
	}
}


<|start_filename|>src/main/java/com/medallia/dsl/interpreter/StreamQueryInterpreter.java<|end_filename|>
package com.medallia.dsl.interpreter;

import com.medallia.data.DataSet;
import com.medallia.data.Segment;
import com.medallia.dsl.Query;
import com.medallia.dsl.ast.Agg;
import com.medallia.dsl.ast.AndExpr;
import com.medallia.dsl.ast.ConstantExpr;
import com.medallia.dsl.ast.Expr;
import com.medallia.dsl.ast.ExprVisitor;
import com.medallia.dsl.ast.InExpr;
import com.medallia.dsl.ast.NotExpr;
import com.medallia.dsl.ast.OrExpr;
import com.medallia.dsl.interpreter.QueryInterpreter.Aggregator;

import java.util.Arrays;

public class StreamQueryInterpreter<T> {

	private final Query<T> query;

	public StreamQueryInterpreter(Query<T> query) {
		this.query = query;
	}

	public T eval(DataSet dataSet) {
		final Expr expr = query.buildExpressionTree();
		final Agg<T> agg = query.aggregateOp.buildAgg();
		final Aggregator aggregator = QueryInterpreter.makeAggregator(dataSet, agg);
		final Filter filter = makeFilter(expr, dataSet);

		final T result = agg.makeResult(dataSet);
		for (Segment segment : dataSet.getSegments()) {
			final int nRows = segment.rawData[0].length;
			for (int row = 0; row < nRows; row++) {
				if (filter.eval(segment, row)) {
					aggregator.process(segment, row, result);
				}
			}
		}
		return result;
	}

	private Filter makeFilter(Expr expression, DataSet dataSet) {
		return expression.visit(new ExprVisitor<Filter>() {
			@Override
			public Filter visit(AndExpr andExpr) {
				Filter left = andExpr.getLeft().visit(this);
				Filter right = andExpr.getRight().visit(this);
				return (segment, row) -> left.eval(segment, row) && right.eval(segment, row);
			}

			@Override
			public Filter visit(InExpr inExpr) {
				final int[] sortedValues = inExpr.getValues().clone();
				Arrays.sort(sortedValues);
				int column = dataSet.getFieldByName(inExpr.getFieldName()).getColumn();
				if (sortedValues.length > 10) {
					return (segment, row) -> Arrays.binarySearch(sortedValues, segment.rawData[column][row]) >= 0;
				} else {
					return (segment, row) -> {
						int val = segment.rawData[column][row];
						for (int sortedValue : sortedValues) {
							if (sortedValue == val) return true;
						}
						return false;
					};
				}
			}

			@Override
			public Filter visit(NotExpr notExpr) {
				Filter target = notExpr.getTarget().visit(this);
				return (segment, row) -> !target.eval(segment, row);
			}

			@Override
			public Filter visit(OrExpr orExpr) {
				Filter left = orExpr.getLeft().visit(this);
				Filter right = orExpr.getRight().visit(this);
				return (segment, row) -> left.eval(segment, row) || right.eval(segment, row);
			}

			@Override
			public Filter visit(ConstantExpr constantExpr) {
				boolean val = constantExpr.value;
				return (segment, row) -> val;
			}
		});
	}

	@FunctionalInterface
	interface Filter {
		boolean eval(Segment segment, int row);
	}
}


<|start_filename|>src/test/java/com/medallia/CrossCheckTest.java<|end_filename|>
package com.medallia;

import com.medallia.data.DataSet;
import com.medallia.data.FieldSpec;
import com.medallia.dsl.FieldStats;
import com.medallia.dsl.Query;
import com.medallia.dsl.compiler.BranchReducingQueryCompiler;
import com.medallia.dsl.compiler.CompiledQueryBase;
import com.medallia.dsl.compiler.QueryCompiler;
import com.medallia.dsl.compiler.unsafe.CCompiledQuery;
import com.medallia.dsl.compiler.unsafe.CQueryCompiler;
import com.medallia.dsl.interpreter.QueryInterpreter;
import com.medallia.dsl.interpreter.StreamQueryInterpreter;
import org.junit.Test;

import static com.medallia.dsl.Aggregate.statsAggregate;
import static com.medallia.dsl.ConditionalExpression.constant;
import static com.medallia.dsl.ConditionalExpression.field;
import static com.medallia.dsl.ConditionalExpression.not;
import static com.medallia.dsl.QueryBuilder.newQuery;
import static org.hamcrest.core.Is.is;
import static org.junit.Assert.assertThat;

/**
 * Simple cross check that all implementations yield the same value for the same query and data
 */
public class CrossCheckTest {
	@Test
	public void test() {
		DataSet dataSet = DataSet.makeRandomDataSet(
				1_000_000, // rows
				50_000,	   // segment size
				new FieldSpec("a", 0, 11),
				new FieldSpec("b", 0, 5),
				new FieldSpec("ltr", 0, 11)
		);

		Query<FieldStats> query = newQuery()
				.filter(
						field("a").in(1, 2, 3)
								.or(field("b").is(3)
								.and(not(constant(false))))
				)
				.filter(not(not(constant(true))))
				.aggregate(statsAggregate("ltr"));


		FieldStats interpreterResult = useQueryInterpreter(dataSet, query);

		FieldStats compiledResult = useSimpleCompiler(dataSet, query);

		FieldStats streamResult = useStreamInterpreter(dataSet, query);

		FieldStats branchResult = useBranchReducingCompiler(dataSet, query);

		FieldStats cResult = useCCompiler(dataSet, query);

		System.out.println("            Interpreted result: " + interpreterResult);
		System.out.println("               Compiled result: " + compiledResult);
		System.out.println("     Stream Interpreted result: " + streamResult);
		System.out.println("Branch reduced compiled result: " + branchResult);
		System.out.println("             C Compiler result: " + cResult);

		assertThat(interpreterResult, is(compiledResult));
		assertThat(compiledResult, is(streamResult));
		assertThat(streamResult, is(branchResult));
		if (cResult != null) {
			assertThat(branchResult, is(cResult));
		}
	}

	private FieldStats useQueryInterpreter(DataSet dataSet, Query<FieldStats> query) {
		final QueryInterpreter<FieldStats> interpreter = new QueryInterpreter<>(query);
		return interpreter.eval(dataSet);
	}

	private FieldStats useSimpleCompiler(DataSet dataSet, Query<FieldStats> query) {
		final QueryCompiler<FieldStats> queryCompiler = new QueryCompiler<>(query, dataSet);
		final CompiledQueryBase<FieldStats> compiledQuery = queryCompiler.compile().get();
		dataSet.getSegments().forEach(compiledQuery::process);
		return compiledQuery.getResult();
	}

	private FieldStats useStreamInterpreter(DataSet dataSet, Query<FieldStats> query) {
		final StreamQueryInterpreter<FieldStats> streamQueryInterpreter = new StreamQueryInterpreter<>(query);
		return streamQueryInterpreter.eval(dataSet);
	}

	private FieldStats useBranchReducingCompiler(DataSet dataSet, Query<FieldStats> query) {
		final BranchReducingQueryCompiler<FieldStats> branchReducingQueryCompiler = new BranchReducingQueryCompiler<>(query, dataSet);
		final CompiledQueryBase<FieldStats> compiledQuery2 = branchReducingQueryCompiler.compile().get();
		dataSet.getSegments().forEach(compiledQuery2::process);
		return compiledQuery2.getResult();
	}

	private FieldStats useCCompiler(DataSet dataSet, Query<FieldStats> query) {
		try {
			final CQueryCompiler<FieldStats> cQueryCompiler = new CQueryCompiler<>(query, dataSet);
			final CompiledQueryBase<FieldStats> compiledQuery3 = cQueryCompiler.compile().get();
			dataSet.getSegments().forEach(compiledQuery3::process);
			return compiledQuery3.getResult();
		} catch (UnsatisfiedLinkError e) {
			// unsafe library not present (most likely), bypass this test
			return null;
		}
	}
}


<|start_filename|>src/main/java/com/medallia/dsl/ast/AggVisitor.java<|end_filename|>
package com.medallia.dsl.ast;

public interface AggVisitor<T> {
	T visit(StatsAggregate statsAggregate);
}


<|start_filename|>src/main/java/com/medallia/dsl/compiler/unsafe/CQueryCompiler.java<|end_filename|>
package com.medallia.dsl.compiler.unsafe;

import com.medallia.codegen.JavaCodeGenerator;
import com.medallia.data.DataSet;
import com.medallia.data.DataSet.FieldDefinition;
import com.medallia.data.FieldSpec;
import com.medallia.data.Segment;
import com.medallia.dsl.FieldStats;
import com.medallia.dsl.Query;
import com.medallia.dsl.ast.Agg;
import com.medallia.dsl.ast.AggVisitor;
import com.medallia.dsl.ast.InExpr;
import com.medallia.dsl.ast.StatsAggregate;
import com.medallia.dsl.compiler.BranchReducingQueryCompiler;
import com.medallia.dsl.compiler.CompiledQueryBase;
import com.medallia.dsl.compiler.QueryCompiler;
import com.medallia.unsafe.Driver;
import com.medallia.unsafe.NativeModule;

import java.io.StringWriter;
import java.util.function.Consumer;
import java.util.function.Supplier;

/**
 * Query compiler that generates C++ code for the query.
 * This one for simplicity is not as generic as the other compilers and expects a {@link FieldStats} as the query result.
 * @param <T> query result type.
 */
public class CQueryCompiler<T extends FieldStats> extends BranchReducingQueryCompiler<T> {

	public CQueryCompiler(Query<T> query, DataSet dataSet) {
		super(query, dataSet);
	}

	@SuppressWarnings("unchecked")
	public Supplier<CompiledQueryBase<T>> compile() {
		final Agg<T> aggregate = query.aggregateOp.buildAgg();
		// This is an abuse, but Java and C/C++ are similar enough that we can us it for both with some care
		final JavaCodeGenerator cg = new JavaCodeGenerator();

		cg.println("#include <jni.h>");

		// Store field ids to amortize lookup time
		cg.println("static jfieldID rawDataFld;");
		cg.println("static jfieldID resultFld;");
		cg.println("static jfieldID sumFld;");
		cg.println("static jfieldID countFld;");

		cg.packageMethod("void", "init").arg("JNIEnv*", "env").arg("jobject", "self")
				.build(methodCg -> {
					methodCg.declare("jclass", "segmentClass", "env->FindClass(\"%s\")", Segment.class.getCanonicalName().replace('.', '/'));
					methodCg.println("rawDataFld = env->GetFieldID(segmentClass, \"rawData\", \"[[I\");");
					methodCg.declare("jclass", "compiledQueryClass", "env->FindClass(\"%s\")", CCompiledQuery.class.getCanonicalName().replace('.', '/'));
					methodCg.println("resultFld = env->GetFieldID(compiledQueryClass, \"result\", \"L" + Object.class.getCanonicalName().replace('.', '/') + ";\");");
					methodCg.declare("jclass", "fieldStatsClass", "env->FindClass(\"%s\")", FieldStats.class.getCanonicalName().replace('.', '/'));
					methodCg.println("sumFld = env->GetFieldID(fieldStatsClass, \"sum\", \"D\");");
					methodCg.println("countFld = env->GetFieldID(fieldStatsClass, \"count\", \"J\");");
				});

		// Generate the process method
		cg.packageMethod("void", "process").arg("JNIEnv*", "env").arg("jobject", "self").arg("jobject", "segment")
				.build(methodCg -> {
					methodCg.declare("jobjectArray", "rawDataObj", "(jobjectArray) env->GetObjectField(segment, rawDataFld)");
					methodCg.declare("jint", "cols", "env->GetArrayLength(rawDataObj)");
					methodCg.println("jint* rawData[cols];");
					methodCg.println("jobject arrays[cols];");
					methodCg.declare("jint", "nRows", "0");
					// Extract column arrays
					methodCg.boundedLoop("i", "0", "(int)cols", loopCg -> {
						loopCg.println("arrays[i] = env->GetObjectArrayElement(rawDataObj, i);");
						loopCg.println("nRows = env->GetArrayLength((jarray)arrays[i]);");
					});
					methodCg.boundedLoop("i", "0", "(int)cols", loopCg -> {
						loopCg.println("rawData[i] = (jint*)env->GetPrimitiveArrayCritical((jarray)arrays[i], 0);");
					});

					// Inner query loop. This assumes we'll be returning a FieldStats object
					methodCg.declare("jdouble", "sum", "0");
					methodCg.declare("jlong", "count", "0");
					methodCg.boundedLoop("row", "0", "(int)nRows", loopCg -> {
						loopCg.ifThen(this::generateFilter, this::generateAggregate);
					});

					// Release column arrays
					methodCg.boundedLoop("i", "0", "(int)cols", loopCg -> {
						loopCg.println("env->ReleasePrimitiveArrayCritical((jarray)arrays[i], rawData[i], 0);");
					});

					// Update result
					methodCg.declare("jobject", "result", "env->GetObjectField(self, resultFld)");
					methodCg.println("env->SetLongField(result, countFld, env->GetLongField(result, countFld) + count);");
					methodCg.println("env->SetDoubleField(result, sumFld, env->GetDoubleField(result, sumFld) + sum);");

				});

		StringWriter sw = new StringWriter();
		cg.generate(sw);
		System.out.println(sw);
		final NativeModule nativeModule = Driver.compileInMemory(sw.toString());
		System.out.println(nativeModule.getErrors());
		// Initialize the native module (setup field ids)
		new CCompiledQuery(nativeModule, null).init();

		return () -> new CCompiledQuery(nativeModule, aggregate.makeResult(dataSet));
	}

	private void generateAggregate(JavaCodeGenerator cg) {
		// This can be done with lambdas, but it's easier to read this way
		Consumer<JavaCodeGenerator> generator = query.aggregateOp.buildAgg().visit(new AggVisitor<Consumer<JavaCodeGenerator>>() {
			@Override
			public Consumer<JavaCodeGenerator> visit(StatsAggregate statsAggregate) {
				return cg -> {
					cg.println(cg.variable("count") + "++;");
					cg.printf("%s += rawData[%d][row];%n", cg.variable("sum"), dataSet.getFieldByName(statsAggregate.getFieldName()).getColumn());
				};
			}

		});

		generator.accept(cg);
	}
}


<|start_filename|>src/main/java/com/medallia/dsl/ast/Expr.java<|end_filename|>
package com.medallia.dsl.ast;

public interface Expr {
	<T> T visit(ExprVisitor<T> visitor);
}


<|start_filename|>src/main/java/com/medallia/dsl/compiler/unsafe/CCompiledQuery.java<|end_filename|>
package com.medallia.dsl.compiler.unsafe;

import com.medallia.data.Segment;
import com.medallia.dsl.compiler.CompiledQueryBase;
import com.medallia.unsafe.Native;
import com.medallia.unsafe.NativeModule;
import com.medallia.unsafe.thunk.NativeBindings;
import com.medallia.unsafe.thunk.ThunkBuilder;

public class CCompiledQuery<T> extends CompiledQueryBase<T> {
	/**
	 * Create native bindings for this class.
	 * Should be kept around for the lifetime of this class.
	 */
	private static final NativeBindings BINDINGS = ThunkBuilder.initializeNative(CCompiledQuery.class);

	/**
	 * Holds pointers to functions needed by the native code.
	 * Must be initialized by a call to {@link com.medallia.unsafe.thunk.NativeBindings#getFunctionPointers(com.medallia.unsafe.NativeModule)}
	 * passing a suitable
	 */
	@Native
	private final long[] functions;

	public CCompiledQuery(NativeModule implementation, T result) {
		functions = BINDINGS.getFunctionPointers(implementation);
		this.result = result;
	}

	public native void init();

	public native void process(Segment segment);
}


<|start_filename|>src/main/java/com/medallia/benchmark/VariantCompiledQueryBenchmark.java<|end_filename|>
package com.medallia.benchmark;

import com.medallia.dsl.FieldStats;
import com.medallia.dsl.compiler.BranchReducingQueryCompiler;
import com.medallia.dsl.compiler.CompiledQueryBase;

import java.util.function.Supplier;

public class VariantCompiledQueryBenchmark
		extends QueryBenchmark<Supplier<CompiledQueryBase<FieldStats>>>
{
	public VariantCompiledQueryBenchmark() {
		super(((dataSet, query) -> new BranchReducingQueryCompiler<>(query, dataSet).compile()), (dataSet, supplier) -> {
			CompiledQueryBase<FieldStats> compiledQuery = supplier.get();
			dataSet.getSegments().forEach(compiledQuery::process);
			return compiledQuery.getResult();
		});
	}
}


<|start_filename|>src/main/java/com/medallia/dsl/compiler/BranchReducingQueryCompiler.java<|end_filename|>
package com.medallia.dsl.compiler;

import com.medallia.codegen.JavaCodeGenerator;
import com.medallia.data.DataSet;
import com.medallia.data.DataSet.FieldDefinition;
import com.medallia.data.FieldSpec;
import com.medallia.dsl.FieldStats;
import com.medallia.dsl.Query;
import com.medallia.dsl.ast.InExpr;

import static com.medallia.dsl.Aggregate.statsAggregate;
import static com.medallia.dsl.ConditionalExpression.field;
import static com.medallia.dsl.QueryBuilder.newQuery;

public class BranchReducingQueryCompiler<T> extends QueryCompiler<T> {

	public BranchReducingQueryCompiler(Query<T> query, DataSet dataSet) {
		super(query, dataSet);
	}

	@Override
	protected String generateInExpr(JavaCodeGenerator cg, InExpr inExpr) {
		final FieldDefinition fieldDef = dataSet.getFieldByName(inExpr.getFieldName());
		final FieldSpec fieldSpec = fieldDef.getFieldSpec();
		final int numValues = fieldSpec.getBound() - fieldSpec.getOrigin();
		final int[] inValues = inExpr.getValues();
		if (inValues.length > 3 && numValues < 64) {
			long mask = 0;
			for (int value : inValues) {
				mask |= 0x8000_0000_0000_0000L >>> (value - fieldSpec.getOrigin());
			}
			if (fieldSpec.getOrigin() == 0) {
				return String.format("(0x%xL << rawData[%d][row]) < 0", mask, fieldDef.getColumn());
			} else {
				return String.format("(0x%xL << (rawData[%d][row]-%dL)) < 0", mask, fieldDef.getColumn(), fieldSpec.getOrigin());
			}
		}
		// Other cases, fall back to the naive one
		return super.generateInExpr(cg, inExpr);
	}

	public static void main(String[] args) {
		DataSet dataSet = DataSet.makeRandomDataSet(
				1_000_000, // rows
				50_000,	   // segment size
				new FieldSpec("a", 0, 11),
				new FieldSpec("b", 0, 5),
				new FieldSpec("ltr", 0, 11)
		);

		Query<FieldStats> query =
				newQuery()
				.filter(field("a").in(1, 2, 3, 9).or(field("b").is(3)))
				.aggregate(statsAggregate("ltr"));

		BranchReducingQueryCompiler<FieldStats> compiler = new BranchReducingQueryCompiler<>(query, dataSet);

		CompiledQueryBase<FieldStats> compiledQuery = compiler.compile().get();

		dataSet.getSegments().forEach(compiledQuery::process);

		System.out.println(compiledQuery.getResult());

	}
}


<|start_filename|>src/main/java/com/medallia/data/FieldSpec.java<|end_filename|>
package com.medallia.data;

public class FieldSpec {
	private final String name;
	private final int origin;
	private final int bound;

	public FieldSpec(String name, int origin, int bound) {
		this.name = name;
		this.origin = origin;
		this.bound = bound;
	}

	public String getName() {
		return name;
	}

	public int getBound() {
		return bound;
	}

	public int getOrigin() {
		return origin;
	}
}


<|start_filename|>src/main/java/com/medallia/dsl/ast/Agg.java<|end_filename|>
package com.medallia.dsl.ast;

import com.medallia.data.DataSet;

public interface Agg<R> {

	R makeResult(DataSet dataSet);

	<T> T visit(AggVisitor<T> visitor);
}


<|start_filename|>src/main/java/com/medallia/dsl/compiler/CompiledQueryBase.java<|end_filename|>
package com.medallia.dsl.compiler;

import com.medallia.data.Segment;

public abstract class CompiledQueryBase<T> {
	protected T result;

	public abstract void process(Segment segment);

	public T getResult() {
		return result;
	}
}
