Predicate.java

package liquidjava.rj_language;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import liquidjava.api.CommandLineLauncher;
import liquidjava.diagnostics.errors.LJError;
import liquidjava.diagnostics.errors.NotFoundError;
import liquidjava.processor.context.AliasWrapper;
import liquidjava.processor.context.Context;
import liquidjava.processor.context.GhostFunction;
import liquidjava.processor.context.GhostState;
import liquidjava.processor.facade.AliasDTO;
import liquidjava.rj_language.ast.BinaryExpression;
import liquidjava.rj_language.ast.Enum;
import liquidjava.rj_language.ast.Expression;
import liquidjava.rj_language.ast.FunctionInvocation;
import liquidjava.rj_language.ast.GroupExpression;
import liquidjava.rj_language.ast.Ite;
import liquidjava.rj_language.ast.LiteralBoolean;
import liquidjava.rj_language.ast.LiteralChar;
import liquidjava.rj_language.ast.LiteralInt;
import liquidjava.rj_language.ast.LiteralLong;
import liquidjava.rj_language.ast.LiteralReal;
import liquidjava.rj_language.ast.UnaryExpression;
import liquidjava.rj_language.ast.Var;
import liquidjava.utils.StaticConstants;
import liquidjava.rj_language.opt.derivation_node.ValDerivationNode;
import liquidjava.rj_language.opt.ExpressionSimplifier;
import liquidjava.rj_language.parsing.RefinementsParser;
import liquidjava.utils.Utils;
import liquidjava.utils.constants.Keys;
import liquidjava.utils.constants.Ops;
import liquidjava.utils.constants.Types;
import spoon.reflect.cu.SourcePosition;
import spoon.reflect.declaration.CtElement;
import spoon.reflect.declaration.CtType;
import spoon.reflect.factory.Factory;

/**
 * Acts as a wrapper for Expression AST
 *
 * @author cgamboa
 */
public class Predicate {

    protected Expression exp;
    protected String prefix;

    /** Create a predicate with the expression true */
    public Predicate() {
        exp = new LiteralBoolean(true);
    }

    /**
     * Create a new predicate with a refinement
     *
     * @param ref
     * @param element
     */
    public Predicate(String ref, CtElement element) throws LJError {
        this(ref, element, element.getParent(CtType.class).getQualifiedName());
    }

    /**
     * Create a new predicate with a refinement and a given prefix for the class
     * 
     * @param ref
     * @param element
     * @param prefix
     */
    public Predicate(String ref, CtElement element, String prefix) throws LJError {
        this.prefix = prefix;
        exp = parse(ref, element);
        if (!(exp instanceof GroupExpression)) {
            exp = new GroupExpression(exp);
        }
        exp = resolveStaticFinalConstants(exp, element);
    }

    /**
     * Walks {@code root}, decorating {@link Enum} nodes that resolve (via reflection, java.lang fallback, or imports
     * declared in {@code context}'s compilation unit) to a {@code static final} primitive/String constant with the
     * corresponding literal expression via {@link Enum#setResolvedLiteral}. The AST shape is preserved, so error
     * messages and counterexamples can render the symbolic {@code Type.CONST} form; the SMT translator emits the
     * literal binding axiom from the decoration. User-defined enums and unresolvable-but-known-type references are left
     * untouched (the SMT side handles them as user enum constants).
     */
    private static Expression resolveStaticFinalConstants(Expression root, CtElement context) throws LJError {
        List<Enum> enums = new ArrayList<>();
        collectEnums(root, enums);
        for (Enum en : enums) {
            Object v = StaticConstants.resolve(en.getTypeName(), en.getConstName(), context);
            Predicate lit = StaticConstants.asLiteralPredicate(v);
            if (lit != null) {
                en.setResolvedLiteral(lit.getExpression());
                continue;
            }
            if (StaticConstants.userTypeExists(en.getTypeName(), context))
                continue; // likely a user-defined enum/class — let SMT translation handle it
            // unresolvable reference — throw an error with a helpful message and import suggestion if possible
            SourcePosition pos = context == null ? null : Utils.getLJAnnotationPosition(context, en.toString());
            String suggested = StaticConstants.findImportCandidate(en.getTypeName(), en.getConstName(), context);
            String hint = suggested != null ? "Add: import " + suggested + ";"
                    : "Add an import for '" + en.getTypeName() + "' if it is a Java class with a static final field";
            String name = en.getTypeName() + "." + en.getConstName();
            NotFoundError error = new NotFoundError(pos, name, "Constant");
            error.setHint(hint);
            throw error;
        }
        return root;
    }

    private static void collectEnums(Expression e, List<Enum> out) {
        if (e instanceof Enum en)
            out.add(en);
        for (Expression c : e.getChildren())
            collectEnums(c, out);
    }

    /**
     * Wrap an already-built expression in a {@link Predicate}. Unlike the string-parsing constructors, this does NOT
     * run static-final-constant resolution. Callers are responsible for ensuring {@code e} contains no
     * {@link liquidjava.rj_language.ast.Enum Enum} nodes that should be resolved to literals; in practice this is fine
     * for AST clones and rewrites that started life from the string constructor.
     */
    public Predicate(Expression e) {
        exp = e;
    }

    protected Expression parse(String ref, CtElement element) throws LJError {
        try {
            return RefinementsParser.createAST(ref, prefix);
        } catch (LJError e) {
            // add location info to error
            if (e.getPosition() == null) {
                SourcePosition pos = Utils.getLJAnnotationPosition(element, ref);
                e.setPosition(pos);
            }
            throw e;
        }
    }

    protected Expression innerParse(String ref, String prefix) throws LJError {
        return RefinementsParser.createAST(ref, prefix);
    }

    public Predicate changeAliasToRefinement(Context context, Factory f) throws LJError {
        Expression ref = getExpression();

        Map<String, AliasDTO> alias = new HashMap<>();
        for (AliasWrapper aw : context.getAliases()) {
            alias.put(aw.getName(), aw.createAliasDTO());
        }

        ref = ref.changeAlias(alias, context, f);
        ref.validateGhostInvocations(context, f);

        return new Predicate(ref);
    }

    public Predicate negate() {
        return new Predicate(new UnaryExpression("!", exp));
    }

    public Predicate substituteVariable(String from, String to) {
        Expression ec = exp.clone();
        ec = ec.substitute(new Var(from), new Var(to));
        return new Predicate(ec);
    }

    public List<String> getVariableNames() {
        List<String> l = new ArrayList<>();
        exp.getVariableNames(l);
        return l;
    }

    public List<GhostState> getStateInvocations(List<GhostState> lgs) {
        if (lgs == null)
            return new ArrayList<>();
        List<String> all = lgs.stream().map(GhostFunction::getQualifiedName).collect(Collectors.toList());
        List<String> toAdd = new ArrayList<>();
        exp.getStateInvocations(toAdd, all);

        List<GhostState> gh = new ArrayList<>();
        for (String n : toAdd) {
            for (GhostState g : lgs)
                if (g.matches(n))
                    gh.add(g);
        }

        return gh;
    }

    /** Change old mentions of previous name to the new name e.g., old(previousName) -> newName */
    public Predicate changeOldMentions(String previousName, String newName) {
        Expression e = exp.clone();
        Expression prev = createVar(previousName).getExpression();
        List<Expression> le = new ArrayList<>();
        le.add(createVar(newName).getExpression());
        e.substituteFunction(Keys.OLD, le, prev);
        return new Predicate(e);
    }

    public Predicate changeStatesToRefinements(List<GhostState> ghostState, String[] toChange) throws LJError {
        Map<String, Expression> nameRefinementMap = new HashMap<>();
        for (GhostState gs : ghostState) {
            if (gs.getRefinement() != null) { // is a state and not a ghost state
                String name = gs.getQualifiedName();
                Expression exp = innerParse(gs.getRefinement().toString(), gs.getPrefix());
                nameRefinementMap.put(name, exp);
                // Also allow simple name lookup to enable hierarchy matching
                String simple = Utils.getSimpleName(name);
                nameRefinementMap.putIfAbsent(simple, exp);
            }
        }
        Expression e = exp.substituteState(nameRefinementMap, toChange);
        return new Predicate(e);
    }

    public Predicate addDerivedStateEqualities(List<GhostState> ghostStates) {
        List<Expression> conjuncts = exp.getConjuncts();
        Map<String, String> stateToInternalState = new HashMap<>();
        for (GhostState gs : ghostStates) {
            if (gs.getRefinement() == null)
                continue;
            Expression ref = gs.getRefinement().getExpression();
            if (ref instanceof GroupExpression ge)
                ref = ge.getExpression();
            if (ref instanceof BinaryExpression be && Ops.EQ.equals(be.getOperator())) {
                FunctionInvocation state = functionInvocation(be.getFirstOperand());
                if (state == null)
                    state = functionInvocation(be.getSecondOperand());
                if (state != null)
                    stateToInternalState.put(gs.getName(), state.getName());
            }
        }

        Predicate result = new Predicate();
        List<Predicate> derivedStates = new ArrayList<>();
        for (Expression conjunct : conjuncts) {
            boolean replaced = false;
            if (conjunct instanceof BinaryExpression be && Ops.EQ.equals(be.getOperator())) {
                FunctionInvocation left = functionInvocation(be.getFirstOperand());
                FunctionInvocation right = functionInvocation(be.getSecondOperand());
                if (left != null && right != null && left.getArgs().size() == 1 && right.getArgs().size() == 1
                        && sameName(left.getName(), right.getName())) {
                    for (Expression stateConjunct : conjuncts) {
                        FunctionInvocation state = functionInvocation(stateConjunct);
                        if (state == null || state.getArgs().size() != 1)
                            continue;
                        for (Map.Entry<String, String> stateName : stateToInternalState.entrySet()) {
                            if (!sameName(stateName.getKey(), state.getName())
                                    || !sameName(stateName.getValue(), left.getName()))
                                continue;

                            Expression known = state.getArgs().get(0);
                            Expression target = known.equals(left.getArgs().get(0)) ? right.getArgs().get(0)
                                    : known.equals(right.getArgs().get(0)) ? left.getArgs().get(0) : null;
                            if (target != null) {
                                derivedStates.add(createInvocation(state.getName(), new Predicate(target)));
                                replaced = true;
                            }
                        }
                    }
                }
            }
            if (!replaced)
                result = createConjunction(result, new Predicate(conjunct));
        }
        for (Predicate derivedState : derivedStates)
            result = createConjunction(result, derivedState);
        return result;
    }

    private static FunctionInvocation functionInvocation(Expression exp) {
        return exp instanceof GroupExpression ge ? functionInvocation(ge.getExpression())
                : exp instanceof FunctionInvocation fi ? fi : null;
    }

    private static boolean sameName(String first, String second) {
        return first.equals(second) || Utils.getSimpleName(first).equals(Utils.getSimpleName(second));
    }

    public boolean isBooleanTrue() {
        return exp.isBooleanTrue();
    }

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

    @Override
    public Predicate clone() {
        return new Predicate(exp.clone());
    }

    @Override
    public boolean equals(Object obj) {
        if (this == obj)
            return true;
        if (obj == null)
            return false;
        if (getClass() != obj.getClass())
            return false;
        Predicate other = (Predicate) obj;
        return exp.equals(other.exp);
    }

    public Expression getExpression() {
        return exp;
    }

    public ValDerivationNode simplify(Context context) {
        // collect aliases from context
        Map<String, AliasDTO> aliases = new HashMap<>();
        for (AliasWrapper aw : context.getAliases()) {
            aliases.put(aw.getName(), aw.createAliasDTO());
        }
        if (CommandLineLauncher.cmdArgs.debugMode) {
            return new ValDerivationNode(exp.clone(), null);
        }
        // simplify expression
        return ExpressionSimplifier.simplify(exp.clone(), aliases);
    }

    private static boolean isBooleanLiteral(Expression expr, boolean value) {
        return expr instanceof LiteralBoolean && expr.isBooleanTrue() == value;
    }

    /**
     * Checks if c2 is a conjunct in c1
     */
    private static boolean containsConjunct(Predicate c1, Predicate c2) {
        if (c1.toString().equals(c2.toString()))
            return true;
        if (c1.getExpression()instanceof BinaryExpression be && be.getOperator().equals(Ops.AND))
            return containsConjunct(new Predicate(be.getFirstOperand()), c2)
                    || containsConjunct(new Predicate(be.getSecondOperand()), c2);
        return false;
    }

    /**
     * Creates a new predicate representing the conjunction of c1 and c2 Contains simplification rules for redundant
     * conjuncts such as not adding conjunct if already present in conjunction
     */
    public static Predicate createConjunction(Predicate c1, Predicate c2) {
        // simplification: (true && x) = x, (false && x) = false
        if (isBooleanLiteral(c1.getExpression(), true))
            return c2;
        if (isBooleanLiteral(c2.getExpression(), true))
            return c1;
        if (isBooleanLiteral(c1.getExpression(), false))
            return c1;
        if (isBooleanLiteral(c2.getExpression(), false))
            return c2;

        // check if conjunct is already present in the conjunction
        if (containsConjunct(c1, c2))
            return c1;
        if (containsConjunct(c2, c1))
            return c2;

        return new Predicate(new BinaryExpression(c1.getExpression(), Ops.AND, c2.getExpression()));
    }

    public static Predicate createDisjunction(Predicate c1, Predicate c2) {
        // simplification: (false || x) = x, (true || x) = true
        if (isBooleanLiteral(c1.getExpression(), false))
            return c2;
        if (isBooleanLiteral(c2.getExpression(), false))
            return c1;
        if (isBooleanLiteral(c1.getExpression(), true))
            return c1;
        if (isBooleanLiteral(c2.getExpression(), true))
            return c2;
        return new Predicate(new BinaryExpression(c1.getExpression(), Ops.OR, c2.getExpression()));
    }

    public static Predicate createEquals(Predicate c1, Predicate c2) {
        return new Predicate(new BinaryExpression(c1.getExpression(), Ops.EQ, c2.getExpression()));
    }

    public static Predicate createITE(Predicate c1, Predicate c2, Predicate c3) {
        return new Predicate(new Ite(c1.getExpression(), c2.getExpression(), c3.getExpression()));
    }

    public static Predicate createLit(String value, String type) {
        Expression exp = switch (type) {
        case Types.BOOLEAN -> new LiteralBoolean(value);
        case Types.INT, Types.SHORT -> new LiteralInt(value);
        case Types.LONG -> new LiteralLong(value);
        case Types.DOUBLE, Types.FLOAT -> new LiteralReal(value);
        case Types.CHAR -> new LiteralChar(value);
        default -> throw new IllegalArgumentException("Unsupported literal type: " + type);
        };
        return new Predicate(exp);
    }

    public static Predicate createOperation(Predicate c1, String op, Predicate c2) {
        return new Predicate(new BinaryExpression(c1.getExpression(), op, c2.getExpression()));
    }

    public static Predicate createVar(String name) {
        return new Predicate(new Var(name));
    }

    public static Predicate createInvocation(String name, Predicate... Predicates) {
        List<Expression> le = new ArrayList<>();
        for (Predicate c : Predicates)
            le.add(c.getExpression());
        return new Predicate(new FunctionInvocation(name, le));
    }
}