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