VariableResolver.java

package liquidjava.rj_language.opt;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import liquidjava.rj_language.ast.BinaryExpression;
import liquidjava.rj_language.ast.Expression;
import liquidjava.rj_language.ast.FunctionInvocation;
import liquidjava.rj_language.ast.Var;

public class VariableResolver {

    /**
     * Extracts variables with constant values from an expression
     * 
     * @param exp
     * 
     * @returns map from variable names to their values
     */
    public static Map<String, Expression> resolve(Expression exp) {
        Map<String, Expression> map = new HashMap<>();

        // extract variable equalities recursively
        resolveRecursive(exp, map);

        // remove variables that were not used in the expression
        map.entrySet().removeIf(entry -> !hasUsage(exp, entry.getKey()));

        // transitively resolve variables
        return resolveTransitive(map);
    }

    /**
     * Recursively extracts variable equalities from an expression (e.g. ... && x == 1 && y == 2 => map: x -> 1, y -> 2)
     * 
     * @param exp
     * @param map
     */
    private static void resolveRecursive(Expression exp, Map<String, Expression> map) {
        if (!(exp instanceof BinaryExpression be))
            return;

        String op = be.getOperator();
        if ("&&".equals(op)) {
            resolveRecursive(be.getFirstOperand(), map);
            resolveRecursive(be.getSecondOperand(), map);
            return;
        }
        if (!"==".equals(op))
            return;

        Expression left = be.getFirstOperand();
        Expression right = be.getSecondOperand();
        String leftKey = substitutionKey(left);
        String rightKey = substitutionKey(right);

        if (leftKey != null && right.isLiteral()) {
            map.put(leftKey, right.clone());
        } else if (rightKey != null && left.isLiteral()) {
            map.put(rightKey, left.clone());
        } else if (left instanceof Var leftVar && right instanceof Var rightVar) {
            // to substitute internal variable with user-facing variable
            if (isInternal(leftVar) && !isInternal(rightVar) && !isReturnVar(leftVar)) {
                map.put(leftVar.getName(), right.clone());
            } else if (isInternal(rightVar) && !isInternal(leftVar) && !isReturnVar(rightVar)) {
                map.put(rightVar.getName(), left.clone());
            } else if (isInternal(leftVar) && isInternal(rightVar)) {
                // to substitute the lower-counter variable with the higher-counter one
                boolean isLeftCounterLower = getCounter(leftVar) <= getCounter(rightVar);
                Var lowerVar = isLeftCounterLower ? leftVar : rightVar;
                Var higherVar = isLeftCounterLower ? rightVar : leftVar;
                if (!isReturnVar(lowerVar) && !isFreshVar(higherVar))
                    map.putIfAbsent(lowerVar.getName(), higherVar.clone());
            }
        } else if (left instanceof Var var && canSubstitute(var, right)) {
            map.put(var.getName(), right.clone());
        } else if (left instanceof FunctionInvocation && !right.toString().contains(leftKey)) {
            map.put(leftKey, right.clone());
        }
    }

    private static String substitutionKey(Expression exp) {
        if (exp instanceof Var var)
            return var.getName();
        if (exp instanceof FunctionInvocation)
            return exp.toString();
        return null;
    }

    /**
     * Handles transitive variable equalities in the map (e.g. map: x -> y, y -> 1 => map: x -> 1, y -> 1)
     * 
     * @param map
     * 
     * @return new map with resolved values
     */
    private static Map<String, Expression> resolveTransitive(Map<String, Expression> map) {
        Map<String, Expression> result = new HashMap<>();
        for (Map.Entry<String, Expression> entry : map.entrySet()) {
            result.put(entry.getKey(), lookup(entry.getValue(), map, new HashSet<>()));
        }
        return result;
    }

    /**
     * Returns the value of a variable by looking up in the map recursively Uses the seen set to avoid circular
     * references (e.g. x -> y, y -> x) which would cause infinite recursion
     * 
     * @param exp
     * @param map
     * @param seen
     * 
     * @return resolved expression
     */
    private static Expression lookup(Expression exp, Map<String, Expression> map, Set<String> seen) {
        String name = substitutionKey(exp);
        if (name == null)
            return exp;

        if (seen.contains(name))
            return exp; // circular reference

        Expression value = map.get(name);
        if (value == null)
            return exp;

        seen.add(name);
        return lookup(value, map, seen);
    }

    /**
     * Checks if a variable is used in the expression (excluding its own definitions)
     *
     * @param exp
     * @param name
     *
     * @return true if used, false otherwise
     */
    private static boolean hasUsage(Expression exp, String name) {
        // exclude own definitions
        if (exp instanceof BinaryExpression binary && "==".equals(binary.getOperator())) {
            Expression left = binary.getFirstOperand();
            Expression right = binary.getSecondOperand();
            if (left instanceof Var v && v.getName().equals(name)
                    && (right.isLiteral() || (!(right instanceof Var) && canSubstitute(v, right))))
                return false;
            if (left instanceof FunctionInvocation && left.toString().equals(name)
                    && (right.isLiteral() || (!(right instanceof Var) && !right.toString().contains(name))))
                return false;
            if (right instanceof Var v && v.getName().equals(name) && left.isLiteral())
                return false;
            if (right instanceof FunctionInvocation && right.toString().equals(name) && left.isLiteral())
                return false;
        }

        // usage found
        if (exp instanceof Var var && var.getName().equals(name)) {
            return true;
        }
        if (exp instanceof FunctionInvocation && exp.toString().equals(name)) {
            return true;
        }

        // recurse children
        if (exp.hasChildren()) {
            for (Expression child : exp.getChildren())
                if (hasUsage(child, name))
                    return true;
        }

        // usage not found
        return false;
    }

    private static int getCounter(Var var) {
        if (!isInternal(var))
            throw new IllegalStateException("Cannot get counter of non-internal variable");
        int lastUnderscore = var.getName().lastIndexOf('_');
        return Integer.parseInt(var.getName().substring(lastUnderscore + 1));
    }

    private static boolean isInternal(Var var) {
        return var.getName().startsWith("#");
    }

    private static boolean isReturnVar(Var var) {
        return var.getName().startsWith("#ret_");
    }

    private static boolean isFreshVar(Var var) {
        return var.getName().startsWith("#fresh_");
    }

    private static boolean canSubstitute(Var var, Expression value) {
        return !isReturnVar(var) && !isFreshVar(var) && !containsVariable(value, var.getName());
    }

    private static boolean containsVariable(Expression exp, String name) {
        if (exp instanceof Var var)
            return var.getName().equals(name);

        if (!exp.hasChildren())
            return false;

        for (Expression child : exp.getChildren()) {
            if (containsVariable(child, name))
                return true;
        }
        return false;
    }
}