001package org.clafer.ir.compiler;
002
003import java.util.Map;
004import org.clafer.collection.Either;
005import org.clafer.ir.IrBoolConstant;
006import org.clafer.ir.IrBoolVar;
007import org.clafer.ir.IrIntConstant;
008import org.clafer.ir.IrIntVar;
009import org.clafer.ir.IrSetConstant;
010import org.clafer.ir.IrSetVar;
011import solver.variables.BoolVar;
012import solver.variables.IntVar;
013import solver.variables.SetVar;
014
015/**
016 * Maps IR non-constant variables to their translated Choco variables. IR
017 * variables can be optimized away. The get<Type>Var methods will return
018 * either the constant it optimized to or the Choco variable. The
019 * get<Type>Value will return the value, regardless of the optimizations.
020 * The get<Type>Value methods are undefined if the solver has not found a
021 * solution yet.
022 *
023 * @author jimmy
024 */
025public class IrSolutionMap {
026
027    private final Map<IrIntVar, IrIntVar> coalescedIntVars;
028    private final Map<IrIntVar, IntVar> intVars;
029    private final Map<IrSetVar, IrSetVar> coalescedSetVars;
030    private final Map<IrSetVar, SetVar> setVars;
031
032    IrSolutionMap(
033            Map<IrIntVar, IrIntVar> coalescedIntVars,
034            Map<IrIntVar, IntVar> intVars,
035            Map<IrSetVar, IrSetVar> coalescedSetVars,
036            Map<IrSetVar, SetVar> setVars) {
037        this.coalescedIntVars = coalescedIntVars;
038        this.intVars = intVars;
039        this.coalescedSetVars = coalescedSetVars;
040        this.setVars = setVars;
041    }
042
043    public Either<Boolean, BoolVar> getBoolVar(IrBoolVar var) {
044        IrBoolVar boolVar = (IrBoolVar) coalescedIntVars.get(var);
045        if (boolVar == null) {
046            boolVar = var;
047        }
048        if (boolVar instanceof IrBoolConstant) {
049            return Either.left(((IrBoolConstant) boolVar).getValue());
050        }
051        return Either.right((BoolVar) intVars.get(boolVar));
052    }
053
054    public Either<Boolean, BoolVar>[] getBoolVars(IrBoolVar... vars) {
055        @SuppressWarnings("unchecked")
056        Either<Boolean, BoolVar>[] bvars = new Either[vars.length];
057        for (int i = 0; i < bvars.length; i++) {
058            bvars[i] = getBoolVar(vars[i]);
059        }
060        return bvars;
061    }
062
063    public boolean getBoolValue(IrBoolVar var) {
064        Either<Boolean, BoolVar> boolVar = getBoolVar(var);
065        return boolVar.isLeft()
066                ? boolVar.getLeft()
067                : boolVar.getRight().getValue() != 0;
068    }
069
070    public boolean[] getBoolValues(IrBoolVar... vars) {
071        boolean[] bvalues = new boolean[vars.length];
072        for (int i = 0; i < bvalues.length; i++) {
073            bvalues[i] = getBoolValue(vars[i]);
074        }
075        return bvalues;
076    }
077
078    public Either<Integer, IntVar> getIntVar(IrIntVar var) {
079        IrIntVar intVar = coalescedIntVars.get(var);
080        if (intVar == null) {
081            intVar = var;
082        }
083        if (intVar instanceof IrIntConstant) {
084            return Either.left(((IrIntConstant) intVar).getValue());
085        }
086        if (intVar instanceof IrBoolConstant) {
087            return Either.left(((IrBoolConstant) intVar).getValue() ? 1 : 0);
088        }
089        return Either.right(intVars.get(intVar));
090    }
091
092    public Either<Integer, IntVar>[] getIntVars(IrIntVar... vars) {
093        @SuppressWarnings("unchecked")
094        Either<Integer, IntVar>[] ivars = new Either[vars.length];
095        for (int i = 0; i < ivars.length; i++) {
096            ivars[i] = getIntVar(vars[i]);
097        }
098        return ivars;
099    }
100
101    public IntVar[] getIntVars() {
102        return intVars.values().toArray(new IntVar[intVars.size()]);
103    }
104
105    public int getIntValue(IrIntVar var) {
106        Either<Integer, IntVar> intVar = getIntVar(var);
107        return intVar.isLeft()
108                ? intVar.getLeft()
109                : intVar.getRight().getValue();
110    }
111
112    public int[] getIntValues(IrIntVar... vars) {
113        int[] ivalues = new int[vars.length];
114        for (int i = 0; i < ivalues.length; i++) {
115            ivalues[i] = getIntValue(vars[i]);
116        }
117        return ivalues;
118    }
119
120    public Either<int[], SetVar> getSetVar(IrSetVar var) {
121        IrSetVar setVar = coalescedSetVars.get(var);
122        if (setVar == null) {
123            setVar = var;
124        }
125        if (setVar instanceof IrSetConstant) {
126            return Either.left(((IrSetConstant) setVar).getValue());
127        }
128        return Either.right(setVars.get(setVar));
129    }
130
131    public Either<int[], SetVar>[] getSetVars(IrSetVar... vars) {
132        @SuppressWarnings("unchecked")
133        Either<int[], SetVar>[] svars = new Either[vars.length];
134        for (int i = 0; i < svars.length; i++) {
135            svars[i] = getSetVar(vars[i]);
136        }
137        return svars;
138    }
139
140    public SetVar[] getSetVars() {
141        return setVars.values().toArray(new SetVar[setVars.size()]);
142    }
143
144    public int[] getSetValue(IrSetVar var) {
145        Either<int[], SetVar> setVar = getSetVar(var);
146        return setVar.isLeft()
147                ? setVar.getLeft()
148                : setVar.getRight().getValue();
149    }
150
151    public int[][] getSetValues(IrSetVar... vars) {
152        int[][] svalues = new int[vars.length][];
153        for (int i = 0; i < svalues.length; i++) {
154            svalues[i] = getSetValue(vars[i]);
155        }
156        return svalues;
157    }
158}