ExprCompiler.java

package com.uber.playground;

import lombok.Value;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Stack;

interface Expr<T extends Expr<T>> {

  ExprClass<T> cls();

  T self();

  default T plus(T rhs) {
    return cls().add(self(), rhs);
  }
}

interface ExprClass<T extends Expr<T>> {

  T literal(int v);

  T add(T lhs, T rhs);
}

interface ExprProgram {
  <T extends Expr<T>> T with(ExprClass<T> cls);
}

final class DataExpr implements ExprClass<DataExpr.Repr> {

  @Override
  public Repr literal(int v) {
    return new Literal(v);
  }

  @Override
  public Repr add(Repr lhs, Repr rhs) {
    return new Add(lhs, rhs);
  }

  public abstract class Repr implements Expr<Repr> {

    abstract Repr simplify();

    abstract <T extends Expr<T>> T as(ExprClass<T> cls);

    @Override
    public ExprClass<Repr> cls() {
      return DataExpr.this;
    }

    @Override
    public Repr self() {
      return this;
    }
  }

  private static boolean isZero(Repr repr) {
    return repr instanceof Literal && ((Literal) repr).getValue() == 0;
  }

  @Value
  public class Literal extends Repr {
    int value;

    @Override
    Repr simplify() {
      return this;
    }

    @Override
    <T extends Expr<T>> T as(ExprClass<T> cls) {
      return cls.literal(value);
    }
  }

  @Value
  public class Add extends Repr {
    Repr lhs, rhs;

    @Override
    Repr simplify() {
      if (isZero(lhs)) {
        return rhs.simplify();
      } else if (isZero(rhs)) {
        return lhs.simplify();
      } else {
        return this;
      }
    }

    @Override
    <T extends Expr<T>> T as(ExprClass<T> cls) {
      return cls.add(lhs.as(cls), rhs.as(cls));
    }
  }
}

abstract class Bytecode {

  public abstract <R> R accept(Visitor<R> visitor);

  interface Visitor<R> {
    R visit(ConstI32 instr);
    R visit(Add instr);
  }

  @Value
  public static class ConstI32 extends Bytecode {
    int value;

    @Override
    public <R> R accept(Visitor<R> visitor) {
      return visitor.visit(this);
    }
  }

  @Value
  public static class Add extends Bytecode {
    @Override
    public <R> R accept(Visitor<R> visitor) {
      return visitor.visit(this);
    }
  }
}

final class BytecodeExpr implements ExprClass<BytecodeExpr.Repr> {

  @Override
  public Repr literal(int v) {
    return new Repr(Arrays.asList(new Bytecode.ConstI32(v)));
  }

  @Override
  public Repr add(Repr lhs, Repr rhs) {
    ArrayList<Bytecode> is = new ArrayList<>();
    is.addAll(lhs.instructions);
    is.addAll(rhs.instructions);
    is.add(new Bytecode.Add());
    return new Repr(is);
  }

  @Value
  public class Repr implements Expr<Repr> {

    List<Bytecode> instructions;

    @Override
    public ExprClass<Repr> cls() {
      return BytecodeExpr.this;
    }

    @Override
    public Repr self() {
      return this;
    }
  }
}

class Compiler {
  static List<Bytecode> compile(ExprProgram program) {
    DataExpr.Repr optimized = program.with(new DataExpr()).simplify();
    System.out.println(optimized);
    BytecodeExpr.Repr compiled = optimized.as(new BytecodeExpr());
    return compiled.getInstructions();
  }
}

class Runtime {
  static int execute(List<Bytecode> is) {
    Stack<Integer> stack = new Stack<>();

    for (Bytecode i : is) {
      if (i instanceof Bytecode.ConstI32) {
        stack.push(((Bytecode.ConstI32) i).getValue());
      } else if (i instanceof Bytecode.Add) {
        Integer x = stack.pop();
        Integer y = stack.pop();
        stack.push(x + y);
      }
    }

    return stack.pop();
  }
}

final class EvalExpr implements ExprClass<EvalExpr.Repr> {

  @Override
  public Repr literal(int v) {
    return new Repr(v);
  }

  @Override
  public Repr add(Repr lhs, Repr rhs) {
    return new Repr(lhs.value + rhs.value);
  }

  @Value
  public class Repr implements Expr<Repr> {

    int value;

    @Override
    public ExprClass<Repr> cls() {
      return EvalExpr.this;
    }

    @Override
    public Repr self() {
      return this;
    }
  }
}

public class ExprCompiler {

  static final ExprProgram p1 = new ExprProgram() {
    @Override
    public <T extends Expr<T>> T with(ExprClass<T> cls) {
      return cls.literal(0).plus(cls.literal(10).plus(cls.literal(5)));
    }
  };

  public static void main(String[] args) {
    System.out.println(p1.with(new EvalExpr()).getValue());
    System.out.println(Runtime.execute(Compiler.compile(p1)));
  }
}