Untyped lambda calculus in Java 21 + preview

JLambda.java

import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.Stream;

sealed interface Expr permits Var, Lam, App {
    default Set<String> freeVars() {
        return switch (this) {
            case Var(var name) -> Set.of(name);
            case Lam(var bound, var body) -> {
                var vars = new HashSet<>(body.freeVars());
                vars.remove(bound);
                yield vars;
            }
            case App(var fn, var arg) ->
                    Stream.concat(fn.freeVars().stream(), arg.freeVars().stream()).collect(Collectors.toSet());
        };
    }

    default Expr subst(String var, Expr to) {
        return switch (this) {
            case Var(var name) when name.equals(var) -> to;
            case Var _ -> this;
            case App(var fn, var arg) -> new App(fn.subst(var, to), arg.subst(var, to));
            case Lam(var bound, _) when bound.equals(var) -> this;
            case Lam(var bound, var body) -> {
                var toFV = to.freeVars();
                if (!toFV.contains(bound)) {
                    yield new Lam(bound, body.subst(var, to));
                } else {
                    var bodyFV = body.freeVars();
                    var fresh = bound;
                    while (toFV.contains(fresh) || bodyFV.contains(fresh)) {
                        fresh = STR."\{fresh}'";
                    }

                    var alphaConvBody = body.subst(bound, new Var(fresh));
                    yield new Lam(fresh, alphaConvBody.subst(var, to));
                }
            }
        };
    }

    default Expr normalize() {
        return switch (this) {
            case Var _ -> this;
            case Lam(var bound, var body) -> new Lam(bound, body.normalize());
            case App(var fn, var arg) -> {
                var normFn = fn.normalize();
                var normArg = arg.normalize();
                yield switch (normFn) {
                    case Lam(var bound, var body) -> {
                        Expr substBody = body.subst(bound, normArg);
                        yield substBody.normalize();
                    }
                    default -> new App(normFn, normArg);
                };
            }
        };
    }
}

record Var(String var) implements Expr {
    static Var of(String name) {
        return new Var(name);
    }

    @Override
    public String toString() {
        return this.var;
    }
}

record Lam(String bound, Expr body) implements Expr {
    static Expr of(List<String> boundVars, Expr body) {
        for (var bv : boundVars.reversed()) {
            body = new Lam(bv, body);
        }
        return body;
    }

    @Override
    public String toString() {
        return STR."λ\{bound}. \{body}";
    }
}

record App(Expr fn, Expr arg) implements Expr {
    static Expr of(Expr fn, Expr... args) {
        for (var arg : args) {
            fn = new App(fn, arg);
        }
        return fn;
    }

    @Override
    public String toString() {
        var fnOut = (fn instanceof Var) ? fn : new Parens<>(fn);
        var argOut = (arg instanceof Var) ? arg : new Parens<>(arg);
        return STR."\{fnOut} \{argOut}";
    }

    record Parens<T>(T inner) {
        @Override
        public String toString() {
            return STR."(\{inner})";
        }
    }
}

// \s. \z. z
Expr zero = Lam.of(List.of("s", "z"), Var.of("z"));

// \n. \s. \z. s (n s z)
Expr succ = Lam.of(
        List.of("n", "s", "z"),
        App.of(
                Var.of("s"),
                App.of(Var.of("n"), Var.of("s"), Var.of("z"))));

// \m. \n. \s. \z. m s (n s z)
Expr plus = Lam.of(
        List.of("m", "n", "s", "z"),
        App.of(Var.of("m"),
                Var.of("s"),
                App.of(Var.of("n"), Var.of("s"), Var.of("z"))));

void main() {
    System.out.println("λ >");
    System.out.println(STR."0 = \{zero}");
    System.out.println(STR."succ = \{succ}");
    System.out.println(STR."plus = \{plus}");

    var one = new App(succ, zero).normalize();
    System.out.println(STR."succ 0 = 1 = \{one}");

    var two = new App(succ, one).normalize();
    System.out.println(STR."succ 1 = 2 = \{two}");

    var three = new App(succ, two).normalize();
    System.out.println(STR."succ 2 = 3 = \{three}");

    var five = new App(new App(plus, two), three).normalize();
    System.out.println(STR."2 + 3 = 5 = \{five}");

    var substEx = new Lam("x", new Var("y")).subst("y", new Var("x"));
    System.out.println(STR."(λx.y)[y:=x] = \{substEx}");
}

/*
λ >
0 = λs. λz. z
succ = λn. λs. λz. s ((n s) z)
plus = λm. λn. λs. λz. (m s) ((n s) z)
succ 0 = 1 = λs. λz. s z
succ 1 = 2 = λs. λz. s (s z)
succ 2 = 3 = λs. λz. s (s (s z))
2 + 3 = 5 = λs. λz. s (s (s (s (s z))))
(λx.y)[y:=x] = λx'. x
*/