1 package jdk.incubator.code.analysis;
2
3 import jdk.incubator.code.Block;
4 import jdk.incubator.code.CodeElement;
5 import jdk.incubator.code.CopyContext;
6 import jdk.incubator.code.Op;
7 import jdk.incubator.code.OpTransformer;
8 import jdk.incubator.code.Value;
9 import jdk.incubator.code.op.CoreOp;
10 import java.util.ArrayList;
11 import java.util.Collections;
12 import java.util.HashMap;
13 import java.util.HashSet;
14 import java.util.LinkedHashSet;
15 import java.util.List;
16 import java.util.Map;
17 import java.util.Objects;
18 import java.util.SequencedSet;
19 import java.util.Set;
20
21 /**
22 * This is an implementation of SSA construction based on
23 * <a href="https://doi.org/10.1007/978-3-642-37051-9">
24 * Simple end Efficient Construction of Static Single Assignment Form (pp 102-122)
25 * </a>.
26 * <p>
27 * This implementation contains some adaptions, notably:
28 * <ul>
29 * <li>Adapt to block parameters rather than phi functions.</li>
30 * <li>Adapt to work with multiple bodies.</li>
31 * </ul>
32 */
33 public class SSAConstruction implements OpTransformer {
34
35 private final Map<CoreOp.VarOp, Map<Block, Val>> currentDef = new HashMap<>();
36 private final Set<Block> sealedBlocks = new HashSet<>();
37 private final Map<Block, Map<CoreOp.VarOp, Phi>> incompletePhis = new HashMap<>();
38
39 // according to the algorithm:
40 // "As only filled blocks may have successors, predecessors are always filled."
41 // In consequence, this means that only filled predecessors should be considered
42 // when recursively searching for a definition
43 private final Map<Block, SequencedSet<Block>> predecessors = new HashMap<>();
44 // as we can't modify the graph at the same time as we analyze it,
45 // we need to store which load op needs to remapped to which value
46 private final Map<CoreOp.VarAccessOp.VarLoadOp, Val> loads = new HashMap<>();
47 private final Map<Block, List<Phi>> additionalParameters = new HashMap<>();
48 // as we look up definitions during the actual transformation again,
49 // we might encounter deleted phis.
50 // we use this set to be able to correct that during transformation
51 private final Set<Phi> deletedPhis = new HashSet<>();
52
53 public static <O extends Op> O transform(O nestedOp) {
54 SSAConstruction construction = new SSAConstruction();
55 construction.prepare(nestedOp);
56 @SuppressWarnings("unchecked")
57 O temp = (O) nestedOp.transform(CopyContext.create(), construction);
58 return temp;
59 }
60
61 private SSAConstruction() {
62 }
63
64 private void prepare(Op nestedOp) {
65 nestedOp.traverse(null, CodeElement.opVisitor((_, op) -> {
66 switch (op) {
67 case CoreOp.VarAccessOp.VarLoadOp load -> {
68 Val val = readVariable(load.varOp(), load.parentBlock());
69 registerLoad(load, val);
70 }
71 case CoreOp.VarAccessOp.VarStoreOp store ->
72 writeVariable(store.varOp(), store.parentBlock(), new Holder(store.storeOperand()));
73 case CoreOp.VarOp initialStore -> {
74 Val val = initialStore.isUninitialized()
75 ? Uninitialized.VALUE
76 : new Holder(initialStore.initOperand());
77 writeVariable(initialStore, initialStore.parentBlock(), val);
78 }
79 case Op.Terminating _ -> {
80 Block block = op.parentBlock();
81 // handle the sealing, i.e. only now make this block a predecessor of its successors
82 for (Block.Reference successor : block.successors()) {
83 Block successorBlock = successor.targetBlock();
84 Set<Block> blocks = this.predecessors.computeIfAbsent(successorBlock, _ -> new LinkedHashSet<>());
85 blocks.add(block);
86 // if this was the last predecessor added to successorBlock, seal it
87 if (blocks.size() == successorBlock.predecessors().size()) {
88 sealBlock(successorBlock);
89 }
90 }
91 }
92 default -> {
93 }
94 }
95 return null;
96 }));
97 }
98
99 private void registerLoad(CoreOp.VarAccessOp.VarLoadOp load, Val val) {
100 this.loads.put(load, val);
101 if (val instanceof Phi phi) {
102 phi.users.add(load);
103 }
104 }
105
106 private void writeVariable(CoreOp.VarOp variable, Block block, Val value) {
107 this.currentDef.computeIfAbsent(variable, _ -> new HashMap<>()).put(block, value);
108 }
109
110 private Val readVariable(CoreOp.VarOp variable, Block block) {
111 Val value = this.currentDef.getOrDefault(variable, Map.of()).get(block);
112 if (value == null
113 // deleted Phi, this is an old reference
114 // due to our 2-step variant of the original algorithm, we might encounter outdated definitions
115 // when we read to prepare block arguments
116 || value instanceof Phi phi && this.deletedPhis.contains(phi)) {
117 return readVariableRecursive(variable, block);
118 }
119 return value;
120 }
121
122 private Val readVariableRecursive(CoreOp.VarOp variable, Block block) {
123 Val value;
124 if (!block.isEntryBlock() && !this.sealedBlocks.contains(block)) {
125 Phi phi = new Phi(variable, block);
126 value = phi;
127 this.incompletePhis.computeIfAbsent(block, _ -> new HashMap<>()).put(variable, phi);
128 this.additionalParameters.computeIfAbsent(block, _ -> new ArrayList<>()).add(phi);
129 } else if (block.isEntryBlock() && variable.ancestorBody() != block.parentBody()) {
130 // we are in an entry block but didn't find a definition yet
131 Block enclosingBlock = block.parent().parent().parent();
132 assert enclosingBlock != null : "def not found in entry block, with no enclosing block";
133 value = readVariable(variable, enclosingBlock);
134 } else if (this.predecessors.get(block).size() == 1) {
135 value = readVariable(variable, this.predecessors.get(block).getFirst());
136 } else {
137 Phi param = new Phi(variable, block);
138 writeVariable(variable, block, param);
139 value = addPhiOperands(variable, param);
140 // To go from Phis to block parameters, we remember that we produced a Phi here.
141 // This means that edges to this block need to pass a value via parameter
142 if (value == param) {
143 this.additionalParameters.computeIfAbsent(block, _ -> new ArrayList<>()).add(param);
144 }
145 }
146 writeVariable(variable, block, value); // cache value for this variable + block
147 return value;
148 }
149
150 private Val addPhiOperands(CoreOp.VarOp variable, Phi value) {
151 for (Block pred : this.predecessors.getOrDefault(value.block(), Collections.emptySortedSet())) {
152 value.appendOperand(readVariable(variable, pred));
153 }
154 return tryRemoveTrivialPhi(value);
155 }
156
157 private Val tryRemoveTrivialPhi(Phi phi) {
158 Val same = null;
159 for (Val op : phi.operands()) {
160 if (op == same || op == phi) {
161 continue;
162 }
163 if (same != null) {
164 return phi;
165 }
166 same = op;
167 }
168 // we shouldn't have phis without operands (other than itself)
169 assert same != null : "phi without different operands";
170 List<Phi> phiUsers = phi.replaceBy(same, this);
171 List<Phi> phis = this.additionalParameters.get(phi.block());
172 if (phis != null) {
173 phis.remove(phi);
174 }
175 for (Phi user : phiUsers) {
176 tryRemoveTrivialPhi(user);
177 }
178 return same;
179 }
180
181 private void sealBlock(Block block) {
182 this.incompletePhis.getOrDefault(block, Map.of()).forEach(this::addPhiOperands);
183 this.sealedBlocks.add(block);
184 }
185
186 // only used during transformation
187
188 private Value resolveValue(CopyContext context, Val val) {
189 return switch (val) {
190 case Uninitialized _ -> throw new IllegalStateException("Uninitialized variable");
191 case Holder holder -> context.getValueOrDefault(holder.value(), holder.value());
192 case Phi phi -> {
193 List<Phi> phis = this.additionalParameters.get(phi.block());
194 int additionalParameterIndex = phis.indexOf(phi);
195 assert additionalParameterIndex >= 0 : "phi not in parameters " + phi;
196 int index = additionalParameterIndex + phi.block().parameters().size();
197 Block.Builder b = context.getBlock(phi.block());
198 yield b.parameters().get(index);
199 }
200 };
201 }
202
203 @Override
204 public Block.Builder apply(Block.Builder block, Op op) {
205 Block originalBlock = op.parentBlock();
206 CopyContext context = block.context();
207 switch (op) {
208 case CoreOp.VarAccessOp.VarLoadOp load -> {
209 Val val = this.loads.get(load);
210 context.mapValue(load.result(), resolveValue(context, val));
211 }
212 case CoreOp.VarOp _, CoreOp.VarAccessOp.VarStoreOp _ -> {
213 }
214 case Op.Terminating _ -> {
215 // make sure outgoing branches are corrected
216 for (Block.Reference successor : originalBlock.successors()) {
217 Block successorBlock = successor.targetBlock();
218 List<Phi> successorParams = this.additionalParameters.getOrDefault(successorBlock, List.of());
219 List<Value> additionalParams = successorParams.stream()
220 .map(phi -> readVariable(phi.variable, originalBlock))
221 .map(val -> resolveValue(context, val))
222 .toList();
223 List<Value> values = context.getValues(successor.arguments());
224 values.addAll(additionalParams);
225 Block.Builder successorBlockBuilder = context.getBlock(successorBlock);
226 context.mapSuccessor(successor, successorBlockBuilder.successor(values));
227 }
228 block.op(op);
229 }
230 default -> block.op(op);
231 }
232 return block;
233 }
234
235 @Override
236 public void apply(Block.Builder block, Block b) {
237 // add the required additional parameters to this block
238 boolean isEntry = b.isEntryBlock();
239 for (Phi phi : this.additionalParameters.getOrDefault(b, List.of())) {
240 if (isEntry) {
241 // Phis in entry blocks denote captured values. Do not add as param but make sure
242 // the original value is used
243 assert phi.operands().size() == 1 : "entry block phi with multiple operands";
244 CopyContext context = block.context();
245 context.mapValue(resolveValue(context, phi), resolveValue(context, phi.operands().getFirst()));
246 } else {
247 block.parameter(phi.variable.varValueType());
248 }
249 }
250
251 // actually visit ops in this block
252 OpTransformer.super.apply(block, b);
253 }
254
255 sealed interface Val {
256 }
257
258 record Holder(Value value) implements Val {
259 }
260
261 enum Uninitialized implements Val {
262 VALUE;
263 }
264
265 record Phi(CoreOp.VarOp variable, Block block, List<Val> operands, Set<Object> users) implements Val {
266 Phi(CoreOp.VarOp variable, Block block) {
267 this(variable, block, new ArrayList<>(), new HashSet<>());
268 }
269
270 void appendOperand(Val val) {
271 this.operands.add(val);
272 if (val instanceof Phi phi) { // load op uses are added separately
273 phi.users.add(this);
274 }
275 }
276
277 @Override
278 public boolean equals(Object obj) {
279 return this == obj;
280 }
281
282 @Override
283 public int hashCode() {
284 return Objects.hash(this.variable, this.block);
285 }
286
287 public List<Phi> replaceBy(Val same, SSAConstruction construction) {
288 List<Phi> usingPhis = new ArrayList<>();
289 for (Object user : this.users) {
290 if (user == this) {
291 continue;
292 }
293 if (same instanceof Phi samePhi) {
294 samePhi.users.add(user);
295 }
296 switch (user) {
297 case Phi phi -> {
298 int i = phi.operands.indexOf(this);
299 assert i >= 0 : "use does not have this as operand";
300 phi.operands.set(i, same);
301 usingPhis.add(phi);
302 }
303 case CoreOp.VarAccessOp.VarLoadOp load -> construction.loads.put(load, same);
304 default -> throw new UnsupportedOperationException(user + ":" + user.getClass());
305 }
306 }
307 if (same instanceof Phi samePhi) {
308 samePhi.users.remove(this);
309 }
310 construction.currentDef.get(this.variable).put(this.block, same);
311 construction.deletedPhis.add(this); // we might not replace all current defs, so mark this phi as deleted
312 this.users.clear();
313 return usingPhis;
314 }
315
316 @Override
317 public String toString() {
318 return "Phi[" + variable.varName() + "(" + block.index() + ")," + "operands: " + operands.size() + "}";
319 }
320 }
321 }