Typestate pattern in Java 21

TypestateJava21Demo.java

import java.util.Map;
import java.util.Objects;

public class TypestateJava21Demo {
    // Shared immutable data for all states (avoids duplication).
    record ConnContext(String host) {
        ConnContext { Objects.requireNonNull(host, "host"); }
        String baseUrl() { return "https://" + host; } // shared, always safe
    }

    // Sealed parent: only invariants / always-valid operations.
    sealed interface Conn permits DisconnectedConn, ConnectedConn, AuthConn {
        ConnContext ctx();
        default String host() { return ctx().host(); }
        default String baseUrl() { return ctx().baseUrl(); }

        // Example of an always-valid transition with a fixed return type:
        // disconnecting from any state always yields Disconnected.
        DisconnectedConn disconnect();
    }

    // State: Disconnected
    record DisconnectedConn(ConnContext ctx) implements Conn {
        DisconnectedConn { Objects.requireNonNull(ctx, "ctx"); }

        // State-specific transition: Disconnected -> Connected
        ConnectedConn connect() {
            // Real code: open socket, etc.
            return new ConnectedConn(ctx);
        }

        @Override
        public DisconnectedConn disconnect() {
            return this; // already disconnected
        }
    }

    // State: Connected
    record ConnectedConn(ConnContext ctx) implements Conn {
        ConnectedConn { Objects.requireNonNull(ctx, "ctx"); }

        // State-specific transition: Connected -> Authenticated
        AuthConn login(String username, String password) {
            Objects.requireNonNull(username, "username");
            Objects.requireNonNull(password, "password");
            // Real code: auth handshake. Here we mint a token.
            var token = "token-for-" + username;
            return new AuthConn(ctx, token);
        }

        @Override
        public DisconnectedConn disconnect() {
            // Real code: close socket.
            return new DisconnectedConn(ctx);
        }
    }

    // State: Authenticated
    record AuthConn(ConnContext ctx, String token) implements Conn {
        AuthConn {
            Objects.requireNonNull(ctx, "ctx");
            Objects.requireNonNull(token, "token");
        }

        // State-specific operation: only available when authenticated
        String fetchUserProfile() {
            return "PROFILE(token=" + token + ", host=" + host() + ")";
        }

        @Override
        public DisconnectedConn disconnect() {
            // Real code: revoke token/close socket.
            return new DisconnectedConn(ctx);
        }
    }

    /**
     * Simulates loading a connection from DB/JSON:
     * returns Conn (unknown-at-compile-time state).
     */
    static Conn loadFromStorage(Map<String, String> row) {
        var host = require(row, "host");
        var state = require(row, "state").toUpperCase();
        var ctx = new ConnContext(host);

        return switch (state) {
            case "DISCONNECTED" -> new DisconnectedConn(ctx);
            case "CONNECTED" -> new ConnectedConn(ctx);
            case "AUTH" -> new AuthConn(ctx, require(row, "token"));
            default -> throw new IllegalArgumentException("Unknown state: " + state);
        };
    }

    static String require(Map<String, String> row, String key) {
        var v = row.get(key);
        if (v == null) throw new IllegalArgumentException("Missing key: " + key);
        return v;
    }

    // --- Application logic ---

    static void runWorkflowFromUnknownState(Conn conn) {
        System.out.println("Loaded: host=" + conn.host() + " baseUrl=" + conn.baseUrl());

        // Narrow once, then regain maximum compile-time safety inside each branch.
        switch (conn) {
            case DisconnectedConn d -> {
                // Only DisconnectedConn has connect()
                ConnectedConn c = d.connect();
                AuthConn a = c.login("alice", "secret");
                System.out.println(a.fetchUserProfile());
            }
            case ConnectedConn c -> {
                // Only ConnectedConn has login()
                AuthConn a = c.login("alice", "secret");
                System.out.println(a.fetchUserProfile());
            }
            case AuthConn a -> {
                // Only AuthConn has fetchUserProfile()
                System.out.println(a.fetchUserProfile());
            }
        }
        
        // Always-valid parent-level operation:
        DisconnectedConn after = conn.disconnect();
        System.out.println("After disconnect: " + after.host());
    }

    public static void main(String[] args) {
        // Example 1: from DB/JSON in AUTH state
        Conn fromDb1 = loadFromStorage(Map.of(
                "host", "example.com",
                "state", "AUTH",
                "token", "t-123"
        ));
        runWorkflowFromUnknownState(fromDb1);

        // Example 2: from DB/JSON in DISCONNECTED state
        Conn fromDb2 = loadFromStorage(Map.of(
                "host", "example.org",
                "state", "DISCONNECTED"
        ));
        runWorkflowFromUnknownState(fromDb2);

        // ---- Compile-time safety examples (uncomment to see compilation errors) ----

        // Conn any = new DisconnectedConn(new ConnContext("x"));
        // any.connect(); // ERROR: connect() is not on Conn

        // ConnectedConn c = new DisconnectedConn(new ConnContext("x")).connect();
        // c.fetchUserProfile(); // ERROR: fetchUserProfile() is not on ConnectedConn

        // DisconnectedConn d = new DisconnectedConn(new ConnContext("x"));
        // d.login("u", "p"); // ERROR: login() is not on DisconnectedConn
    }
}