TypedExpr.java

package com.uber.playground.typedexpr;

import lombok.Value;
import lombok.val;

import java.util.function.Function;

abstract class Type<T extends Type<T>> {
  private Type() {}

  public static Int int_() {
    return Int.INSTANCE;
  }

  public static <A extends Type<A>, B extends Type<B>> Fun<A, B> fun(Type<A> arg, Type<B> body) {
    return new Fun<>(arg, body);
  }

  @Value
  public static class Int extends Type<Int> {
    private static final Int INSTANCE = new Int();

    @Override
    public String toString() {
      return "Int";
    }
  }

  @Value
  public static class Fun<A extends Type<A>, B extends Type<B>> extends Type<Fun<A, B>> {
    Type<A> arg;
    Type<B> body;

    @Override
    public String toString() {
      return "(" + arg + " -> " + body + ")";
    }
  }
}

interface Expr<T extends Type<T>> {

}

interface ExprClass {

  Expr<Type.Int> literal(int value);

  Expr<Type.Int> add(Expr<Type.Int> lhs, Expr<Type.Int> rhs);

  <A extends Type<A>, B extends Type<B>> Expr<Type.Fun<A, B>> lambda(Type<A> argTy, Function<Expr<A>, Expr<B>> body);

  <A extends Type<A>, B extends Type<B>> Expr<B> apply(Expr<Type.Fun<A, B>> fun, Expr<A> arg);
}

interface ExprProgram {
  Expr<Type.Int> with(ExprClass cls);
}

final class DataExpr implements ExprClass {

  int counter = 0;

  private String freshName() {
    return "x" + counter++;
  }

  @Override
  public Expr<Type.Int> literal(int value) {
    return new Lit(value);
  }

  @Override
  public Expr<Type.Int> add(Expr<Type.Int> lhs, Expr<Type.Int> rhs) {
    return new Add(lhs, rhs);
  }

  @Override
  public <A extends Type<A>, B extends Type<B>> Expr<Type.Fun<A, B>> lambda(Type<A> argTy, Function<Expr<A>, Expr<B>> body) {
    Var<A> var = new Var<>(argTy, freshName());
    return new Lam<>(var, body.apply(var));
  }

  @Override
  public <A extends Type<A>, B extends Type<B>> Expr<B> apply(Expr<Type.Fun<A, B>> fun, Expr<A> arg) {
    return new App<>(fun, arg);
  }

  public static abstract class Repr<T extends Type<T>> implements Expr<T> {}

  @Value
  public static class Lit extends Repr<Type.Int> {
    int value;

    @Override
    public String toString() {
      return "" + value;
    }
  }

  @Value
  public static class Var<T extends Type<T>> extends Repr<T> {
    Type<T> ty;
    String name;

    @Override
    public String toString() {
      return name;
    }
  }

  @Value
  public static class Add extends Repr<Type.Int> {
    Expr<Type.Int> lhs, rhs;

    @Override
    public String toString() {
      return "(" + lhs + " + " + rhs + ")";
    }
  }

  @Value
  public static class Lam<A extends Type<A>, B extends Type<B>> extends Repr<Type.Fun<A, B>> {
    Var<A> arg;
    Expr<B> body;

    @Override
    public String toString() {
      return "(\\" + arg.getName() + ":" + arg.getTy() + " -> " + body + ")";
    }
  }

  @Value
  public static class App<A extends Type<A>, B extends Type<B>> extends Repr<B> {
    Expr<Type.Fun<A, B>> fun;
    Expr<A> arg;

    @Override
    public String toString() {
      return fun + " " + arg;
    }
  }
}

final class EvalExpr implements ExprClass {

  @Override
  public Expr<Type.Int> literal(int value) {
    return new Int(value);
  }

  @Override
  public Expr<Type.Int> add(Expr<Type.Int> lhs, Expr<Type.Int> rhs) {
    Int a = (Int) lhs;
    Int b = (Int) rhs;
    return new Int(a.value + b.value);
  }

  @Override
  public <A extends Type<A>, B extends Type<B>> Expr<Type.Fun<A, B>> lambda(Type<A> argTy, Function<Expr<A>, Expr<B>> body) {
    return new Lam<>(body);
  }

  @Override
  public <A extends Type<A>, B extends Type<B>> Expr<B> apply(Expr<Type.Fun<A, B>> fun, Expr<A> arg) {
    Lam<A, B> lam = (Lam<A, B>) fun;
    return lam.body.apply(arg);
  }

  public static abstract class Val<T extends Type<T>> implements Expr<T> {}

  @Value
  public static class Int extends Val<Type.Int> {
    int value;
  }

  @Value
  public static class Lam<A extends Type<A>, B extends Type<B>> extends Val<Type.Fun<A, B>> {
    Function<Expr<A>, Expr<B>> body;
  }
}

public class TypedExpr {

  static ExprProgram p1 = cls ->
  {
    Expr<Type.Fun<Type.Int, Type.Int>> f = cls.lambda(Type.int_(), x -> cls.add(x, cls.literal(10)));
    return cls.apply(
            f,
            cls.literal(5));
  };

  static ExprProgram p2 = cls ->
  {
    val z = cls.lambda(Type.fun(Type.int_(), Type.fun(Type.int_(), Type.int_())), f -> cls.apply(cls.apply(f, cls.literal(1)), cls.literal(2)));
    return cls.apply(
            z,
            cls.lambda(Type.int_(), x -> cls.lambda(Type.int_(), y -> cls.add(x, y))));
  };

  public static void main(String[] args) {
    System.out.println(p1.with(new DataExpr())); // (\x0:Int -> (x0 + 10)) 5
    System.out.println(p1.with(new EvalExpr())); // EvalExpr.Int(value=15)

    System.out.println(p2.with(new DataExpr())); // (\x0:(Int -> (Int -> Int)) -> x0 1 2) (\x1:Int -> (\x2:Int -> (x1 + x2)))
    System.out.println(p2.with(new EvalExpr())); // EvalExpr.Int(value=3)
  }
}