ASTMatch.java

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

public class ASTMatch implements ASTNode {
        ASTNode exp;
        List<ASTPattern> patterns;
        List<ASTNode> bodies;

        public ASTMatch(ASTNode e, List<ASTPattern> patterns, List<ASTNode> bodies) {
                this.exp = e;
                this.patterns = patterns;
                this.bodies = bodies;
        }

        @Override
        public IValue eval(Environment<IValue> e) throws InterpreterError {
                IValue v = exp.eval(e);

                for (int i = 0; i < patterns.size(); i++) {
                        ASTPattern pat = patterns.get(i);
                        ASTNode body = bodies.get(i);

                        if (v instanceof VStructTuple st && st.getTag().equals(pat.getTag()) && st.getValues().size() >= pat.getPvs().size()) {
                                Environment<IValue> en = e.beginScope();
                                List<String> patVars = pat.getPvs();
                                List<IValue> values = st.getValues();
                                for (int j = 0; j < patVars.size(); j++) {
                                        if (!patVars.get(j).equals("_")) {  // _ is the #Nil() placeholder
                                                en.assoc(patVars.get(j), values.get(j));
                                        }
                                }
                                return body.eval(en);
                        }
                }
                throw new InterpreterError("match: no pattern matched for " + v.toStr());
        }

        @Override
        public ASTType typecheck(Environment<ASTType> e) throws TypeCheckError, InterpreterError {
                ASTType exprType = exp.typecheck(e);
                if (!(exprType instanceof ASTTEnum enumType)) {
                        throw new TypeCheckError("match: expression must be an enum type, got " + exprType.toStr());
                }

                // Checks exhaustiveness
                boolean hasWildcard = patterns.stream().anyMatch(p -> p.getTag() == null || p.getTag().equals("_"));

                if (!hasWildcard) {
                        for (Map.Entry<String, List<ASTType>> variant : enumType.getOptions()) {
                                String tag = variant.getKey();
                                boolean covered = patterns.stream().anyMatch(p -> p.getTag() != null && p.getTag().equals(tag));
                                if (!covered)
                                        throw new TypeCheckError("match: non-exhaustive patterns, missing tag " + tag);
                        }
                }

                ASTType resultType = null;

                for (int i = 0; i < patterns.size(); i++) {
                        ASTPattern pat = patterns.get(i);
                        ASTNode body = bodies.get(i);
                        String tag = pat.getTag();
                        List<String> patVars = pat.getPvs();
                        boolean isWildcard = (tag == null || tag.equals("_"));
                        List<ASTType> fieldTypes = new ArrayList<>();

                        if (!isWildcard) {
                                fieldTypes = enumType.getOptions().stream().filter(entry -> entry.getKey().equals(tag))
                                        .map(Map.Entry::getValue).findFirst()
                                        .orElseThrow(() -> new TypeCheckError("match: tag " + tag + " not in enum"));

                                boolean isNilCase = patVars.size() == 1 && patVars.get(0).equals("_") && fieldTypes.isEmpty();

                                if (!isNilCase && patVars.size() > fieldTypes.size()) {
                                        throw new TypeCheckError("match: pattern has more variables than enum fields for " + tag);
                                }
                        }

                        Environment<ASTType> patEnv = e.beginScope();
                        
                        if (!isWildcard) {
                                for (int j = 0; j < patVars.size(); j++) {
                                        if (!patVars.get(j).equals("_")) {
                                                patEnv.assoc(patVars.get(j), fieldTypes.get(j));
                                        }
                                }
                        }

                        ASTType branchType = body.typecheck(patEnv);
                        patEnv.endScope();

                        branchType = ASTTypeDef.unfold(branchType, e);

                        if (resultType == null) {
                                resultType = branchType;
                        } else {
                                resultType = resultType.computeLUB(branchType, e);
                        }
                }
                return ASTTypeDef.unfold(resultType, e);
        }

        @Override
        public boolean isFreeOutsideLambda(String id) {
                if (exp.isFreeOutsideLambda(id)) return true;
                for (ASTNode body : bodies)
                        if (body.isFreeOutsideLambda(id))
                                return true;
                return false;
        }
}