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.traverse(null, CodeElement.opVisitor((_, op) -> {
 85             switch (op) {
 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.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             return null;
115         }));
116     }
117 
118     private void registerLoad(CoreOp.VarAccessOp.VarLoadOp load, Val val) {
119         this.loads.put(load, val);
120         if (val instanceof Phi phi) {
121             phi.users.add(load);
122         }
123     }
124 
125     private void writeVariable(CoreOp.VarOp variable, Block block, Val value) {
126         this.currentDef.computeIfAbsent(variable, _ -> new HashMap<>()).put(block, value);
127     }
128 
129     private Val readVariable(CoreOp.VarOp variable, Block block) {
130         Val value = this.currentDef.getOrDefault(variable, Map.of()).get(block);
131         if (value == null
132             // deleted Phi, this is an old reference
133             // due to our 2-step variant of the original algorithm, we might encounter outdated definitions
134             // when we read to prepare block arguments
135             || value instanceof Phi phi && this.deletedPhis.contains(phi)) {
136             return readVariableRecursive(variable, block);
137         }
138         return value;
139     }
140 
141     private Val readVariableRecursive(CoreOp.VarOp variable, Block block) {
142         Val value;
143         if (!block.isEntryBlock() && !this.sealedBlocks.contains(block)) {
144             Phi phi = new Phi(variable, block);
145             value = phi;
146             this.incompletePhis.computeIfAbsent(block, _ -> new HashMap<>()).put(variable, phi);
147             this.additionalParameters.computeIfAbsent(block, _ -> new ArrayList<>()).add(phi);
148         } else if (block.isEntryBlock() && variable.ancestorBody() != block.ancestorBody()) {
149             // we are in an entry block but didn't find a definition yet
150             Block enclosingBlock = block.parent().parent().parent();
151             assert enclosingBlock != null : "def not found in entry block, with no enclosing block";
152             value = readVariable(variable, enclosingBlock);
153         } else if (this.predecessors.get(block).size() == 1) {
154             value = readVariable(variable, this.predecessors.get(block).getFirst());
155         } else {
156             Phi param = new Phi(variable, block);
157             writeVariable(variable, block, param);
158             value = addPhiOperands(variable, param);
159             // To go from Phis to block parameters, we remember that we produced a Phi here.
160             // This means that edges to this block need to pass a value via parameter
161             if (value == param) {
162                 this.additionalParameters.computeIfAbsent(block, _ -> new ArrayList<>()).add(param);
163             }
164         }
165         writeVariable(variable, block, value); // cache value for this variable + block
166         return value;
167     }
168 
169     private Val addPhiOperands(CoreOp.VarOp variable, Phi value) {
170         for (Block pred : this.predecessors.getOrDefault(value.block(), Collections.emptySortedSet())) {
171             value.appendOperand(readVariable(variable, pred));
172         }
173         return tryRemoveTrivialPhi(value);
174     }
175 
176     private Val tryRemoveTrivialPhi(Phi phi) {
177         Val same = null;
178         for (Val op : phi.operands()) {
179             if (op == same || op == phi) {
180                 continue;
181             }
182             if (same != null) {
183                 return phi;
184             }
185             same = op;
186         }
187         // we shouldn't have phis without operands (other than itself)
188         assert same != null : "phi without different operands";
189         List<Phi> phiUsers = phi.replaceBy(same, this);
190         List<Phi> phis = this.additionalParameters.get(phi.block());
191         if (phis != null) {
192             phis.remove(phi);
193         }
194         for (Phi user : phiUsers) {
195             Val res = tryRemoveTrivialPhi(user);
196             if (same == user) {
197                 same = res;
198             }
199         }
200         return same;
201     }
202 
203     private void sealBlock(Block block) {
204         this.incompletePhis.getOrDefault(block, Map.of()).forEach(this::addPhiOperands);
205         this.sealedBlocks.add(block);
206     }
207 
208     // only used during transformation
209 
210     private Value resolveValue(CodeContext context, Val val) {
211         return switch (val) {
212             case Uninitialized _ -> throw new IllegalStateException("Uninitialized variable");
213             case Holder holder -> context.getValueOrDefault(holder.value(), holder.value());
214             case Phi phi -> {
215                 List<Phi> phis = this.additionalParameters.get(phi.block());
216                 int additionalParameterIndex = phis.indexOf(phi);
217                 assert additionalParameterIndex >= 0 : "phi not in parameters " + phi;
218                 int index = additionalParameterIndex + phi.block().parameters().size();
219                 Block.Builder b = context.getBlock(phi.block());
220                 yield b.parameters().get(index);
221             }
222         };
223     }
224 
225     @Override
226     public Block.Builder acceptOp(Block.Builder block, Op op) {
227         Block originalBlock = op.ancestorBlock();
228         CodeContext context = block.context();
229         switch (op) {
230             case CoreOp.VarAccessOp.VarLoadOp load -> {
231                 Val val = this.loads.get(load);
232                 context.mapValue(load.result(), resolveValue(context, val));
233             }
234             case CoreOp.VarOp _, CoreOp.VarAccessOp.VarStoreOp _ -> {
235             }
236             case Op.Terminating _ -> {
237                 // make sure outgoing branches are corrected
238                 for (Block.Reference successor : originalBlock.successors()) {
239                     Block successorBlock = successor.targetBlock();
240                     List<Phi> successorParams = this.additionalParameters.getOrDefault(successorBlock, List.of());
241                     List<Value> additionalParams = successorParams.stream()
242                             .map(phi -> readVariable(phi.variable, originalBlock))
243                             .map(val -> resolveValue(context, val))
244                             .toList();
245                     List<Value> values = context.getValues(successor.arguments());
246                     values.addAll(additionalParams);
247                     Block.Builder successorBlockBuilder = context.getBlock(successorBlock);
248                     context.mapSuccessor(successor, successorBlockBuilder.successor(values));
249                 }
250                 block.op(op);
251             }
252             default -> block.op(op);
253         }
254         return block;
255     }
256 
257     @Override
258     public void acceptBlock(Block.Builder block, Block b) {
259         // add the required additional parameters to this block
260         boolean isEntry = b.isEntryBlock();
261         for (Phi phi : this.additionalParameters.getOrDefault(b, List.of())) {
262             if (isEntry) {
263                 // Phis in entry blocks denote captured values. Do not add as param but make sure
264                 // the original value is used
265                 assert phi.operands().size() == 1 : "entry block phi with multiple operands";
266                 CodeContext context = block.context();
267                 context.mapValue(resolveValue(context, phi), resolveValue(context, phi.operands().getFirst()));
268             } else {
269                 block.parameter(phi.variable.varValueType());
270             }
271         }
272 
273         // actually visit ops in this block
274         CodeTransformer.super.acceptBlock(block, b);
275     }
276 
277     sealed interface Val {
278     }
279 
280     record Holder(Value value) implements Val {
281     }
282 
283     enum Uninitialized implements Val {
284         VALUE;
285     }
286 
287     record Phi(CoreOp.VarOp variable, Block block, List<Val> operands, Set<Object> users) implements Val {
288         Phi(CoreOp.VarOp variable, Block block) {
289             this(variable, block, new ArrayList<>(), new HashSet<>());
290         }
291 
292         void appendOperand(Val val) {
293             this.operands.add(val);
294             if (val instanceof Phi phi) { // load op uses are added separately
295                 phi.users.add(this);
296             }
297         }
298 
299         @Override
300         public boolean equals(Object obj) {
301             return this == obj;
302         }
303 
304         @Override
305         public int hashCode() {
306             return Objects.hash(this.variable, this.block);
307         }
308 
309         public List<Phi> replaceBy(Val same, SSABraun construction) {
310             List<Phi> usingPhis = new ArrayList<>();
311             for (Object user : this.users) {
312                 if (user == this) {
313                     continue;
314                 }
315                 if (same instanceof Phi samePhi) {
316                     samePhi.users.add(user);
317                 }
318                 switch (user) {
319                     case Phi phi -> {
320                         int i = phi.operands.indexOf(this);
321                         assert i >= 0 : "use does not have this as operand";
322                         phi.operands.set(i, same);
323                         usingPhis.add(phi);
324                     }
325                     case CoreOp.VarAccessOp.VarLoadOp load -> construction.loads.put(load, same);
326                     default -> throw new UnsupportedOperationException(user + ":" + user.getClass());
327                 }
328             }
329             if (same instanceof Phi samePhi) {
330                 samePhi.users.remove(this);
331             }
332             construction.currentDef.get(this.variable).put(this.block, same);
333             construction.deletedPhis.add(this); // we might not replace all current defs, so mark this phi as deleted
334             this.users.clear();
335             return usingPhis;
336         }
337 
338         @Override
339         public String toString() {
340             return "Phi[" + variable.varName() + "(" + block.index() + ")," + "operands: " + operands.size() + "}";
341         }
342     }
343 }