StateMachineParser.java

package liquidjava.fsm;

import java.net.URI;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Comparator;
import java.util.List;

import liquidjava.rj_language.ast.*;
import liquidjava.rj_language.parsing.RefinementsParser;
import liquidjava.processor.VCImplication;
import liquidjava.rj_language.ast.formatter.ExpressionFormatter;
import liquidjava.rj_language.ast.formatter.VariableFormatter;
import liquidjava.utils.Utils;
import spoon.Launcher;
import spoon.reflect.CtModel;
import spoon.reflect.declaration.CtAnnotation;
import spoon.reflect.declaration.CtClass;
import spoon.reflect.declaration.CtConstructor;
import spoon.reflect.declaration.CtElement;
import spoon.reflect.declaration.CtExecutable;
import spoon.reflect.declaration.CtMethod;
import spoon.reflect.declaration.CtType;
import spoon.reflect.cu.SourcePosition;

public class StateMachineParser {

    private static final String STATE_SET_ANNOTATION = "StateSet";
    private static final String STATE_REFINEMENT_ANNOTATION = "StateRefinement";
    private static final String REFINEMENT_ANNOTATION = "Refinement";
    private static final String EXTERNAL_REFINEMENTS_FOR_ANNOTATION = "ExternalRefinementsFor";

    private record TransitionSource(String from, String cond) {}

    /**
     * Parses a class or interface for the given uri and extracts the state machine information
     * @param uri the document URI
     * @return StateMachine or null if none found
     */
    public static StateMachine parse(String uri) {
        return parse(uri, null, null, null);
    }

    public static StateMachine parseWithErrorContext(String uri, SourcePosition declarationPosition,
            Expression expectedState, VCImplication foundState) {
        return parse(uri, declarationPosition, expectedState, foundState);
    }

    private static StateMachine parse(String uri, SourcePosition declarationPosition,
            Expression expectedState, VCImplication foundState) {
        try {
            String filePath = Paths.get(new URI(uri)).toString();
            Launcher launcher = new Launcher();
            launcher.getEnvironment().setNoClasspath(true);
            launcher.getEnvironment().setAutoImports(true);
            launcher.addInputResource(filePath);
            launcher.buildModel();
            CtModel model = launcher.getModel();

            // get class or interface
            CtType<?> ctType = getType(model);
            if (ctType == null)
                return null; // no class or interface found

            // extract class name and states
            List<String> states = getStates(ctType);
            if (states == null || states.isEmpty())
                return null; // no states found
            String className = getClassName(ctType);
            String simpleClassName = Utils.getSimpleName(className);

            // get initial transitions and method transitions
            List<StateMachineInitialTransition> factoryTransitions = getFactoryTransitions(ctType, simpleClassName, states);
            List<StateMachineInitialTransition> initialTransitions = getInitialTransitions(ctType, simpleClassName, states);
            // interfaces only have an implicit initial state when there are no factory methods
            if (initialTransitions.isEmpty() && (ctType instanceof CtClass<?> || factoryTransitions.isEmpty())) {
                initialTransitions = getDefaultInitialTransitions(ctType, simpleClassName, states.get(0));
            }
            initialTransitions = new ArrayList<>(initialTransitions);
            initialTransitions.addAll(factoryTransitions);
            List<StateMachineTransition> transitions = getTransitions(ctType, simpleClassName, states);
            if (transitions.isEmpty())
                return null; // no transitions found

            StateMachineErrorContext errorContext = declarationPosition == null && expectedState == null
                    && foundState == null ? null : getErrorContext(ctType, simpleClassName, states,
                    declarationPosition, expectedState, foundState);
            return new StateMachine(className, states, transitions, initialTransitions, errorContext);

        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }

    private static StateMachineErrorContext getErrorContext(CtType<?> ctType, String className, List<String> states,
            SourcePosition declarationPosition, Expression expectedState, VCImplication foundState) {
        String receiver = getStateReceiver(expectedState, states);
        String receiverName = receiver == null ? null : VariableFormatter.withoutInstance(receiver);
        List<Expression> foundPremises = getReceiverPremises(foundState, receiver);
        return new StateMachineErrorContext(
                getCalledMethod(ctType, className, declarationPosition),
                getActualStates(expectedState, foundState, states),
                receiverName,
                getExpectedStates(expectedState, states, receiver),
                expectedState == null ? null : formatCondition(expectedState, receiver, receiverName),
                getFoundCondition(foundState, foundPremises, receiver, receiverName),
                getDefiniteState(foundPremises, states, receiver));
    }

    /**
     * Gets the name of the method, constructor or static factory whose declaration contains the given position
     * The verifier reports the position of the state refinement annotation value, so the position is matched by range
     * Constructors are named after the class
     */
    private static String getCalledMethod(CtType<?> ctType, String className, SourcePosition declarationPosition) {
        if (declarationPosition == null || !declarationPosition.isValidPosition()) return null;
        List<CtExecutable<?>> executables = new ArrayList<>(ctType.getMethods());
        if (ctType instanceof CtClass<?> ctClass) {
            executables.addAll(ctClass.getConstructors());
        }
        return executables.stream()
                .filter(executable -> contains(executable.getPosition(), declarationPosition))
                .min(Comparator.comparingInt(executable -> executable.getPosition().getSourceEnd()
                        - executable.getPosition().getSourceStart()))
                .map(executable -> executable instanceof CtConstructor<?> ? className : executable.getSimpleName())
                .orElse(null);
    }

    private static boolean contains(SourcePosition outer, SourcePosition inner) {
        return outer.isValidPosition()
                && outer.getFile() != null && inner.getFile() != null
                && outer.getFile().toPath().toAbsolutePath().normalize()
                        .equals(inner.getFile().toPath().toAbsolutePath().normalize())
                && outer.getSourceStart() <= inner.getSourceStart()
                && inner.getSourceEnd() <= outer.getSourceEnd();
    }

    private static List<String> getActualStates(Expression expectedState, VCImplication foundState,
            List<String> states) {
        String receiver = getStateReceiver(expectedState, states);
        if (receiver == null || foundState == null) return List.of();

        List<String> foundStates = new ArrayList<>();
        for (VCImplication premise = foundState; premise != null; premise = premise.getNext()) {
            if (premise.hasBinder() && receiver.equals(premise.getName()) && premise.getRefinement() != null) {
                foundStates.addAll(getStateExpressions(premise.getRefinement().getExpression(), states, receiver));
            }
        }
        return states.stream().filter(foundStates::contains).toList();
    }

    /**
     * Gets the states accepted by the expected condition for the receiver, in state set order
     */
    private static List<String> getExpectedStates(Expression expectedState, List<String> states, String receiver) {
        if (expectedState == null || receiver == null) return List.of();
        List<String> expectedStates = getStateExpressions(expectedState, states, receiver);
        return states.stream().filter(expectedStates::contains).toList();
    }

    /**
     * Gets the found premises about the receiver: those bound to it and unbound ones that mention it
     */
    private static List<Expression> getReceiverPremises(VCImplication foundState, String receiver) {
        List<Expression> premises = new ArrayList<>();
        if (receiver == null) return premises;
        for (VCImplication premise = foundState; premise != null; premise = premise.getNext()) {
            if (premise.getRefinement() == null) continue;
            Expression expr = premise.getRefinement().getExpression();
            boolean bound = premise.hasBinder() && receiver.equals(premise.getName());
            boolean mentioned = !premise.hasBinder() && premise.getRefinement().getVariableNames().contains(receiver);
            if ((bound || mentioned) && !expr.isBooleanTrue()) {
                premises.add(expr);
            }
        }
        return premises;
    }

    private static String getFoundCondition(VCImplication foundState, List<Expression> foundPremises,
            String receiver, String receiverName) {
        if (foundState == null) return null;
        if (receiver == null) return ExpressionFormatter.format(foundState.toPredicate().getExpression());
        return formatCondition(joinConjunction(foundPremises), receiver, receiverName);
    }

    private static Expression joinConjunction(List<Expression> exprs) {
        Expression conjunction = null;
        for (Expression expr : exprs) {
            conjunction = conjunction == null ? expr : new BinaryExpression(conjunction, "&&", expr);
        }
        return conjunction == null ? new LiteralBoolean(true) : conjunction;
    }

    /**
     * Formats a condition for display, showing the receiver by its source name (e.g. root instead of root¹⁷)
     */
    private static String formatCondition(Expression expr, String receiver, String receiverName) {
        if (receiver != null && receiverName != null) {
            expr = expr.substitute(new Var(receiver), new Var(receiverName));
        }
        return ExpressionFormatter.format(expr);
    }

    /**
     * Gets the single state the found premises establish for the receiver, or null if there is none
     * A state is established only when it is a top-level conjunct; any other conjunct about the receiver's
     * states (disjunction, conditional, ...) makes the result null, except negations of other states.
     * Negations alone never establish a state since the verifier does not bound the state set
     */
    private static String getDefiniteState(List<Expression> premises, List<String> states, String receiver) {
        if (receiver == null) return null;
        List<String> found = new ArrayList<>();
        List<String> excluded = new ArrayList<>();
        for (Expression premise : premises) {
            for (Expression conjunct : premise.getConjuncts()) {
                String state = getReceiverState(conjunct, states, receiver);
                String negated = conjunct instanceof UnaryExpression unary && unary.getOp().equals("!")
                        ? getReceiverState(unary.getExpression(), states, receiver) : null;
                if (state != null) {
                    if (!found.contains(state)) found.add(state);
                } else if (negated != null) {
                    excluded.add(negated);
                } else if (mentionsReceiverState(conjunct, states, receiver)) {
                    return null;
                }
            }
        }
        if (found.size() != 1 || excluded.contains(found.get(0))) return null;
        return found.get(0);
    }

    private static String getReceiverState(Expression expr, List<String> states, String receiver) {
        return expr instanceof FunctionInvocation ? getStateName(expr, states, receiver) : null;
    }

    private static boolean mentionsReceiverState(Expression expr, List<String> states, String receiver) {
        if (getReceiverState(expr, states, receiver) != null) return true;
        return expr.getChildren().stream().anyMatch(child -> mentionsReceiverState(child, states, receiver));
    }

    private static String getStateReceiver(Expression expr, List<String> states) {
        if (expr == null) return null;
        if (expr instanceof FunctionInvocation invocation
                && findState(invocation.getName(), states) != null
                && !invocation.getArgs().isEmpty()
                && invocation.getArgs().get(0) instanceof Var receiver) {
            return receiver.getName();
        }
        for (Expression child : expr.getChildren()) {
            String receiver = getStateReceiver(child, states);
            if (receiver != null) return receiver;
        }
        return null;
    }

    /**
     * Finds the first class or interface in the model (we assume only one per file)
     * @param model the CtModel
     * @return CtType or null if none found
     */
    private static CtType<?> getType(CtModel model) {
        for (CtType<?> type : model.getAllTypes()) {
            if (type.isClass() || type.isInterface())
                return type;
        }
        return null;
    }

    /**
     * Gets the qualified name from a class or interface
     * Uses name from ExternalRefinementsFor if present, otherwise uses class or interface name
     * @param ctType the CtType (class or interface)
     * @return class name
     */
    private static String getClassName(CtType<?> ctType) {
        for (CtAnnotation<?> annotation : ctType.getAnnotations()) {
            if (annotation.getAnnotationType().getSimpleName().equals(EXTERNAL_REFINEMENTS_FOR_ANNOTATION)) {
                String qualifiedName = (String) annotation.getValueAsObject("value");
                return qualifiedName;
            }
        }
        return ctType.getQualifiedName();
    }

    /**
     * Gets the possible states from a class or interface
     * @param ctType the CtType (class or interface)
     * @return list of states
     */
    private static List<String> getStates(CtType<?> ctType) {
        for (CtAnnotation<?> annotation : ctType.getAnnotations()) {
            if (annotation.getAnnotationType().getSimpleName().equals(STATE_SET_ANNOTATION)) {
                String[] stateArray = (String[]) annotation.getValueAsObject("value");
                return List.of(stateArray);
            }
        }
        return null;
    }

    /**
     * Gets the elements that represent constructors (actual constructors for classes, methods named after the class for interfaces)
     * @param ctType the CtType (class or interface)
     * @param className the class name
     * @return collection of constructor elements
     */
    private static Collection<? extends CtExecutable<?>> getConstructorElements(CtType<?> ctType, String className) {
        if (ctType instanceof CtClass<?> ctClass) {
            return ctClass.getConstructors();
        }
        // for interfaces the constructors are methods with the same name as the class
        return ctType.getMethods().stream().filter(m -> m.getSimpleName().equals(className)).toList();
    }

    /**
     * Gets the initial transitions from a class or interface
     * If not explicitly defined, uses the first state in the state set
     * @param ctType the CtType (class or interface)
     * @param className the class name
     * @param states the list of states
     * @return initial transitions
     */
    private static List<StateMachineInitialTransition> getInitialTransitions(CtType<?> ctType, String className, List<String> states) {
        List<StateMachineInitialTransition> initialTransitions = new ArrayList<>();
        for (CtExecutable<?> element : getConstructorElements(ctType, className)) {
            String signature = getConstructorSignature(element, className);
            boolean hasExplicitInitialState = false;
            for (CtAnnotation<?> annotation : element.getAnnotations()) {
                if (annotation.getAnnotationType().getSimpleName().equals(STATE_REFINEMENT_ANNOTATION)) {
                    String to = annotation.getValueAsString("to");
                    List<StateMachineInitialTransition> transitions = parseInitialTransitions(to, states, signature);
                    initialTransitions.addAll(transitions);
                    hasExplicitInitialState |= !transitions.isEmpty();
                }
            }
            if (!hasExplicitInitialState) {
                initialTransitions.add(new StateMachineInitialTransition(states.get(0), null, signature));
            }
        }
        return initialTransitions;
    }

    private static List<StateMachineInitialTransition> getDefaultInitialTransitions(CtType<?> ctType,
            String className, String firstState) {
        Collection<? extends CtExecutable<?>> constructors = getConstructorElements(ctType, className);
        if (constructors.isEmpty()) {
            String implicitConstructor = ctType instanceof CtClass<?> ? "new " + className + "()" : null;
            return List.of(new StateMachineInitialTransition(firstState, null, implicitConstructor));
        }
        return constructors.stream()
                .map(element -> new StateMachineInitialTransition(firstState, null,
                        getConstructorSignature(element, className)))
                .toList();
    }

    private static String getConstructorSignature(CtExecutable<?> element, String className) {
        return "new " + className + "(" + getParameterTypes(element) + ")";
    }

    /**
     * Gets the initial transitions from static factory methods (methods that return a new object in a given state)
     * For classes these must be static, for interfaces any method returning the class type is considered
     * @param ctType the CtType (class or interface)
     * @param className the class name
     * @param states the list of states
     * @return initial transitions
     */
    private static List<StateMachineInitialTransition> getFactoryTransitions(CtType<?> ctType, String className, List<String> states) {
        List<StateMachineInitialTransition> initialTransitions = new ArrayList<>();
        for (CtMethod<?> method : ctType.getMethods()) {
            if (ctType.isClass() && !method.isStatic())
                continue;
            if (method.getType() == null || !method.getType().getSimpleName().equals(className))
                continue;
            String signature = className + "." + method.getSimpleName() + "(" + getParameterTypes(method) + ")";
            for (CtAnnotation<?> annotation : method.getAnnotations()) {
                if (annotation.getAnnotationType().getSimpleName().equals(REFINEMENT_ANNOTATION)) {
                    String refinement = annotation.getValueAsString("value");
                    initialTransitions.addAll(parseInitialTransitions(refinement, states, signature));
                }
            }
        }
        return initialTransitions;
    }

    private static String getParameterTypes(CtExecutable<?> element) {
        return String.join(", ", element.getParameters().stream()
                .map(parameter -> parameter.getType() == null ? "?" : parameter.getType().getSimpleName())
                .toList());
    }

    /**
     * Gets transitions from a class or interface
     * @param ctType the CtType (class or interface)
     * @param className the class name
     * @param states the list of states
     * @return list of StateMachineTransition
     */
    private static List<StateMachineTransition> getTransitions(CtType<?> ctType, String className, List<String> states) {
        List<StateMachineTransition> transitions = new ArrayList<>();
        for (CtMethod<?> method : ctType.getMethods()) {
            // for interfaces we skip constructor methods (methods with same name as class)
            if (ctType.isInterface() && method.getSimpleName().equals(className))
                continue;
            for (CtAnnotation<?> annotation : method.getAnnotations()) {
                if (annotation.getAnnotationType().getSimpleName().equals(STATE_REFINEMENT_ANNOTATION)) {
                    transitions.addAll(getTransitions(annotation, method.getSimpleName(), states));
                }
            }
        }
        return transitions;
    }

    /**
     * Gets transitions from the given annotation
     * @param ann the CtAnnotation
     * @param method the method name
     * @param states the list of states
     * @return list of StateMachineTransition
     */
    private static List<StateMachineTransition> getTransitions(CtAnnotation<?> ann, String method, List<String> states) {
        List<StateMachineTransition> transitions = new ArrayList<>();
        String from = ann.getValueAsString("from");
        String to = ann.getValueAsString("to");

        // parse from and to expressions
        List<TransitionSource> fromSources = from.isEmpty() ? allStateSources(states) : parsePrecondition(from, states);
        List<TransitionSource> toSources = !from.isEmpty() && to.isEmpty()
                ? parseStateSources(from, states)
                : parsePostcondition(to, states);

        // create transitions for each combination of from and to states
        for (TransitionSource fromSource : fromSources) {
            for (TransitionSource toSource : toSources) {
                transitions.add(new StateMachineTransition(fromSource.from(), toSource.from(), method, fromSource.cond(), toSource.cond()));
            }
        }
        return transitions;
    }

    private static List<TransitionSource> allStateSources(List<String> states) {
        List<TransitionSource> sources = new ArrayList<>();
        for (String state : states) {
            sources.add(new TransitionSource(state, null));
        }
        return sources;
    }

    private static List<TransitionSource> parsePrecondition(String expr, List<String> states) {
        Expression ast = RefinementsParser.createAST(expr, "");
        List<TransitionSource> sources = getTransitionSources(ast, states, false);
        if (sources.isEmpty()) {
            sources = addCondition(allStateSources(states), ast.toString());
        }
        return sources;
    }

    private static List<TransitionSource> parsePostcondition(String expr, List<String> states) {
        if (expr == null || expr.isEmpty()) return new ArrayList<>();
        Expression ast = RefinementsParser.createAST(expr, "");
        return getTransitionSources(ast, states, true);
    }

    private static List<TransitionSource> getTransitionSources(Expression expr, List<String> states, boolean stateOnlyDisjunctions) {
        String state = getStateName(expr, states);
        if (state != null) {
            return List.of(new TransitionSource(state, null));
        }

        if (expr instanceof BinaryExpression bin) {
            String op = bin.getOperator();
            if (op.equals("&&")) {
                return getConjunctionSources(bin, states, stateOnlyDisjunctions);
            } else if (op.equals("||")) {
                return stateOnlyDisjunctions ? getStateSources(bin, states) : getDisjunctionSources(bin, states);
            }
        } else if (expr instanceof UnaryExpression unary) {
            if (unary.getOp().equals("!")) {
                List<String> negatedStates = getStateExpressions(unary.getExpression(), states);
                if (!negatedStates.isEmpty()) {
                    List<TransitionSource> sources = new ArrayList<>();
                    for (String possibleState : states) {
                        if (!negatedStates.contains(possibleState)) {
                            sources.add(new TransitionSource(possibleState, null));
                        }
                    }
                    return sources;
                }
            }
        } else if (expr instanceof Ite ite) {
            List<TransitionSource> sources = new ArrayList<>();
            sources.addAll(addCondition(getTransitionSources(ite.getThen(), states, stateOnlyDisjunctions), ite.getCondition().toString()));
            sources.addAll(addCondition(getTransitionSources(ite.getElse(), states, stateOnlyDisjunctions), negateCondition(ite.getCondition())));
            return sources;
        }
        return new ArrayList<>();
    }

    private static List<TransitionSource> getConjunctionSources(BinaryExpression bin, List<String> states, boolean stateOnlyDisjunctions) {
        List<TransitionSource> leftSources = getTransitionSources(bin.getFirstOperand(), states, stateOnlyDisjunctions);
        List<TransitionSource> rightSources = getTransitionSources(bin.getSecondOperand(), states, stateOnlyDisjunctions);

        if (leftSources.isEmpty() && rightSources.isEmpty()) {
            return new ArrayList<>();
        } else if (leftSources.isEmpty()) {
            return addCondition(rightSources, bin.getFirstOperand().toString());
        } else if (rightSources.isEmpty()) {
            return addCondition(leftSources, bin.getSecondOperand().toString());
        }

        List<TransitionSource> sources = new ArrayList<>(leftSources);
        sources.addAll(rightSources);
        return sources;
    }

    private static List<TransitionSource> getDisjunctionSources(BinaryExpression bin, List<String> states) {
        List<TransitionSource> sources = new ArrayList<>();
        addDisjunctionBranch(sources, bin.getFirstOperand(), states);
        addDisjunctionBranch(sources, bin.getSecondOperand(), states);
        return removeRedundantGuardedSources(sources);
    }

    private static void addDisjunctionBranch(List<TransitionSource> sources, Expression expr, List<String> states) {
        List<TransitionSource> parsedSources = getTransitionSources(expr, states, false);
        if (parsedSources.isEmpty()) {
            sources.addAll(addCondition(allStateSources(states), expr.toString()));
        } else {
            sources.addAll(parsedSources);
        }
    }

    private static List<TransitionSource> removeRedundantGuardedSources(List<TransitionSource> sources) {
        List<TransitionSource> filteredSources = new ArrayList<>();
        for (TransitionSource source : sources) {
            if (source.cond() == null || !hasUnguardedSource(sources, source.from())) {
                filteredSources.add(source);
            }
        }
        return filteredSources;
    }

    private static boolean hasUnguardedSource(List<TransitionSource> sources, String from) {
        for (TransitionSource source : sources) {
            if (source.from().equals(from) && source.cond() == null) {
                return true;
            }
        }
        return false;
    }

    private static List<TransitionSource> addCondition(List<TransitionSource> sources, String cond) {
        List<TransitionSource> guardedSources = new ArrayList<>();
        for (TransitionSource source : sources) {
            guardedSources.add(new TransitionSource(source.from(), combineConditions(source.cond(), cond)));
        }
        return guardedSources;
    }

    private static String combineConditions(String first, String second) {
        if (first == null || first.isEmpty()) {
            return second;
        }
        if (second == null || second.isEmpty()) {
            return first;
        }
        return first + " && " + second;
    }

    private static String negateCondition(Expression expr) {
        String condition = expr.toString();
        if (expr instanceof Var || expr instanceof FunctionInvocation) {
            return "!" + condition;
        }
        return "!(" + condition + ")";
    }

    private static List<StateMachineInitialTransition> parseInitialTransitions(String expr, List<String> states,
            String constructorSignature) {
        if (expr == null || expr.isEmpty()) return new ArrayList<>();
        List<StateMachineInitialTransition> initialTransitions = new ArrayList<>();
        for (TransitionSource source : parsePostcondition(expr, states)) {
            initialTransitions.add(new StateMachineInitialTransition(source.from(), source.cond(), constructorSignature));
        }
        return initialTransitions;
    }

    private static List<TransitionSource> parseStateSources(String expr, List<String> states) {
        if (expr == null || expr.isEmpty()) return new ArrayList<>();
        Expression ast = RefinementsParser.createAST(expr, "");
        return getStateSources(ast, states);
    }

    private static List<TransitionSource> getStateSources(Expression expr, List<String> states) {
        List<TransitionSource> sources = new ArrayList<>();
        for (String state : getStateExpressions(expr, states)) {
            sources.add(new TransitionSource(state, null));
        }
        return sources;
    }

    /**
     * Gets state names from an expression AST recursively
     * @param expr the expression
     * @param states the list of possible states
     * @return list of states
     */
    private static List<String> getStateExpressions(Expression expr, List<String> states) {
        return getStateExpressions(expr, states, null);
    }

    private static List<String> getStateExpressions(Expression expr, List<String> states, String receiver) {
        List<String> stateExpressions = new ArrayList<>();
        String state = getStateName(expr, states, receiver);
        if (state != null) {
            stateExpressions.add(state);
        } else if (expr instanceof BinaryExpression bin) {
            stateExpressions.addAll(getStateExpressions(bin.getFirstOperand(), states, receiver));
            stateExpressions.addAll(getStateExpressions(bin.getSecondOperand(), states, receiver));
        } else if (expr instanceof UnaryExpression unary) {
            if (unary.getOp().equals("!")) {
                // all except those in the expression
                List<String> negatedStates = getStateExpressions(unary.getExpression(), states, receiver);
                if (!negatedStates.isEmpty()) {
                    for (String possibleState : states) {
                        if (!negatedStates.contains(possibleState)) {
                            stateExpressions.add(possibleState);
                        }
                    }
                }
            }
        } else if (expr instanceof Ite ite) {
            // combine states from then and else branches
            stateExpressions.addAll(getStateExpressions(ite.getThen(), states, receiver));
            stateExpressions.addAll(getStateExpressions(ite.getElse(), states, receiver));
        }
        return stateExpressions;
    }

    private static String getStateName(Expression expr, List<String> states) {
        return getStateName(expr, states, null);
    }

    private static String getStateName(Expression expr, List<String> states, String receiver) {
        if (expr instanceof Var var) {
            return findState(var.getName(), states);
        } else if (expr instanceof FunctionInvocation func) {
            String state = findState(func.getName(), states);
            if (state == null || receiver == null) return state;
            if (!func.getArgs().isEmpty()
                    && func.getArgs().get(0) instanceof Var target
                    && receiver.equals(target.getName())) {
                return state;
            }
        }
        return null;
    }

    private static String findState(String name, List<String> states) {
        if (states.contains(name)) {
            return name;
        }
        String simpleName = Utils.getSimpleName(name);
        for (String state : states) {
            if (state.equals(simpleName) || Utils.getSimpleName(state).equals(simpleName)) {
                return state;
            }
        }
        return null;
    }
}