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