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 
 26 package jdk.incubator.code.analysis;
 27 
 28 import jdk.incubator.code.*;
 29 import jdk.incubator.code.dialect.core.CoreOp;
 30 
 31 import java.util.*;
 32 import java.util.stream.Collectors;
 33 
 34 /**
 35  * An implementation of SSA construction based on
 36  * "Efficiently Computing Static Single Assignment Form and the Control Dependence Graph" by Ron Cytron et. al.
 37  * <p>
 38  * This implementation contains some adaptions, notably:
 39  * <ul>
 40  *     <li>Adapt to block parameters rather than phi functions.</li>
 41  *     <li>Adapt to work with multiple bodies.</li>
 42  * </ul>
 43  */
 44 final class SSACytron {
 45     private SSACytron() {
 46     }
 47 
 48     /**
 49      * Applies an SSA transformation to an operation with bodies, replacing operations that declare variables and
 50      * access them with the use of values they depend on or additional block parameters.
 51      * <p>
 52      * The operation should first be in lowered form before applying this transformation.
 53      * <p>
 54      * Note: this implementation does not currently work correctly when a variable is stored to within an exception
 55      * region and read from outside as a result of catching an exception. In such cases a complete transformation may be
 56      * not possible and such variables will need to be retained.
 57      *
 58      * @param nestedOp the operation with bodies
 59      * @return the transformed operation
 60      * @param <T> the operation type
 61      */
 62     static <T extends Op & Op.Nested> T transform(T nestedOp) {
 63         Map<Block, Set<CoreOp.VarOp>> joinPoints = new HashMap<>();
 64         Map<CoreOp.VarAccessOp.VarLoadOp, Object> loadValues = new HashMap<>();
 65         Map<Block.Reference, List<Object>> joinSuccessorValues = new HashMap<>();
 66 
 67         Map<Body, Boolean> visited = new HashMap<>();
 68         Map<Block, Map<CoreOp.VarOp, Block.Parameter>> joinBlockArguments = new HashMap<>();
 69         @SuppressWarnings("unchecked")
 70         T ssaOp = (T) nestedOp.transform(CopyContext.create(), (block, op) -> {
 71             // Compute join points and value mappings for body
 72             visited.computeIfAbsent(op.ancestorBody(), b -> {
 73                 findJoinPoints(b, joinPoints);
 74                 variableToValue(b, joinPoints, loadValues, joinSuccessorValues);
 75                 return true;
 76             });
 77 
 78             if (op instanceof CoreOp.VarOp || op instanceof CoreOp.VarAccessOp) {
 79                 // Drop var operations
 80                 if (op instanceof CoreOp.VarAccessOp.VarLoadOp vl) {
 81                     // Replace result of load
 82                     Object loadValue = loadValues.get(vl);
 83                     CopyContext cc = block.context();
 84                     Value v = loadValue instanceof VarOpBlockArgument vba
 85                             ? joinBlockArguments.get(vba.b()).get(vba.vop())
 86                             : cc.getValue((Value) loadValue);
 87                     cc.mapValue(op.result(), v);
 88                 }
 89             } else if (op instanceof Op.Terminating) {
 90                 for (Block.Reference s : op.successors()) {
 91                     List<Object> joinValues = joinSuccessorValues.get(s);
 92                     // Successor has join values
 93                     if (joinValues != null) {
 94                         CopyContext cc = block.context();
 95 
 96                         // Lazily append target block arguments
 97                         joinBlockArguments.computeIfAbsent(s.targetBlock(), b -> {
 98                             Block.Builder bb = cc.getBlock(b);
 99                             return joinPoints.get(b).stream().collect(Collectors.toMap(
100                                     varOp -> varOp,
101                                     varOp -> bb.parameter(varOp.varValueType())));
102                         });
103 
104                         // Append successor arguments
105                         List<Value> values = new ArrayList<>();
106                         for (Object o : joinValues) {
107                             Value v = o instanceof VarOpBlockArgument vba
108                                     ? joinBlockArguments.get(vba.b()).get(vba.vop())
109                                     : cc.getValue((Value) o);
110                             values.add(v);
111                         }
112 
113                         // Map successor with append arguments
114                         List<Value> toArgs = cc.getValues(s.arguments());
115                         toArgs.addAll(values);
116                         Block.Reference toS = cc.getBlock(s.targetBlock()).successor(toArgs);
117                         cc.mapSuccessor(s, toS);
118                     }
119                 }
120 
121                 block.apply(op);
122             } else {
123                 block.apply(op);
124             }
125 
126             return block;
127         });
128         return ssaOp;
129     }
130 
131     record VarOpBlockArgument(Block b, CoreOp.VarOp vop) {
132     }
133 
134     enum Uninitialized {
135         UNINITIALIZED;
136     }
137 
138     // @@@ Check for var uses in exception regions
139     //     A variable cannot be converted to SAA form if the variable is stored
140     //     to in an exception region and accessed from an associated catch region
141 
142     static void variableToValue(Body body,
143                                 Map<Block, Set<CoreOp.VarOp>> joinPoints,
144                                 Map<CoreOp.VarAccessOp.VarLoadOp, Object> loadValues,
145                                 Map<Block.Reference, List<Object>> joinSuccessorValues) {
146         Map<CoreOp.VarOp, Deque<Object>> variableStack = new HashMap<>();
147         Node top = buildDomTree(body.entryBlock(), body.immediateDominators());
148         variableToValue(top, variableStack, joinPoints, loadValues, joinSuccessorValues);
149     }
150 
151     /**
152      * Replaces usages of a variable with the corresponding value, from a given block node in the dominator tree.
153      * <p>
154      * The result of a {@code VarLoadOp} for variable, {@code V} say the result of a {@code VarOp} operation,
155      * is replaced with the value passed as an operand to the immediately dominating {@code VarStoreOp} that operates
156      * on {@code V}, or a block argument representing the equivalent of a phi-value of {@code V}.
157      * After which, any related {@code VarOp}, {@code VarLoadOp}, or {@code VarStoreOp} operations are removed.
158      *
159      * @param n             the node in the dominator tree
160      * @param variableStack the variable stack
161      * @param joinPoints    the join points
162      * @implNote See "Efficiently Computing Static Single Assignment Form and the Control Dependence Graph" by Ron Cytron et. al.
163      * Section 5.2 and Figure 12.
164      */
165     static void variableToValue(Node n,
166                                 Map<CoreOp.VarOp, Deque<Object>> variableStack,
167                                 Map<Block, Set<CoreOp.VarOp>> joinPoints,
168                                 Map<CoreOp.VarAccessOp.VarLoadOp, Object> loadValues,
169                                 Map<Block.Reference, List<Object>> joinSuccessorValues) {
170         int size = n.b().ops().size();
171 
172         // Check if V is associated with block argument (phi)
173         // Push argument onto V's stack
174         {
175             Set<CoreOp.VarOp> varOps = joinPoints.get(n.b());
176             if (varOps != null) {
177                 varOps.forEach(v -> {
178                     assert variableStack.containsKey(v);
179                     variableStack.get(v).push(new VarOpBlockArgument(n.b(), v));
180                 });
181             }
182         }
183 
184         {
185             for (int i = 0; i < size - 1; i++) {
186                 Op op = n.b().ops().get(i);
187 
188                 if (op instanceof CoreOp.VarOp varOp) {
189                     // Initial value assigned to variable
190                     Object current = varOp.isUninitialized()
191                             ? Uninitialized.UNINITIALIZED
192                             : op.operands().get(0);
193                     assert !variableStack.containsKey(varOp);
194                     variableStack.computeIfAbsent(varOp, _ -> new ArrayDeque<>())
195                             .push(current);
196                 } else if (op instanceof CoreOp.VarAccessOp.VarStoreOp storeOp) {
197                     // Value assigned to variable
198                     Value current = op.operands().get(1);
199                     variableStack.get(storeOp.varOp()).push(current);
200                 } else if (op instanceof CoreOp.VarAccessOp.VarLoadOp loadOp &&
201                            loadOp.varOp().ancestorBody() == op.ancestorBody()) {
202                     Object to = peekAtCurrentVariable(variableStack, loadOp.varOp());
203                     loadValues.put(loadOp, to);
204                 } else if (op instanceof Op.Nested) {
205                     // Traverse descendant variable loads for variables
206                     // declared in the block's parent body
207                     op.traverse(null, (_, ce) -> {
208                         if (ce instanceof CoreOp.VarAccessOp.VarLoadOp loadOp &&
209                                 loadOp.varOp().ancestorBody() == op.ancestorBody()) {
210                             Object to = peekAtCurrentVariable(variableStack, loadOp.varOp());
211                             loadValues.put(loadOp, to);
212                         }
213                         return null;
214                     });
215                 }
216             }
217 
218             // Add successor args for joint points
219             for (Block.Reference succ : n.b().successors()) {
220                 Set<CoreOp.VarOp> varOps = joinPoints.get(succ.targetBlock());
221                 if (varOps != null) {
222                     List<Object> joinValues = varOps.stream()
223                             .map(vop -> peekAtCurrentVariable(variableStack, vop)).toList();
224                     joinSuccessorValues.put(succ, joinValues);
225                 }
226             }
227 
228             // The result of a VarOp, a variable value, can only be used in VarStoreOp and VarLoadOp
229             // therefore there is no need to check existing successor arguments
230         }
231 
232         // Traverse children of dom tree
233         for (Node y : n.children()) {
234             variableToValue(y, variableStack, joinPoints, loadValues, joinSuccessorValues);
235         }
236 
237         // Pop off values for variables
238         {
239             Set<CoreOp.VarOp> varOps = joinPoints.get(n.b());
240             if (varOps != null) {
241                 varOps.forEach(v -> {
242                     variableStack.get(v).pop();
243                 });
244             }
245 
246             for (int i = 0; i < size - 1; i++) {
247                 Op op = n.b().ops().get(i);
248 
249                 if (op instanceof CoreOp.VarOp varOp) {
250                     variableStack.get(varOp).pop();
251                 } else if (op instanceof CoreOp.VarAccessOp.VarStoreOp storeOp) {
252                     variableStack.get(storeOp.varOp()).pop();
253                 }
254             }
255         }
256     }
257 
258     static Object peekAtCurrentVariable(Map<CoreOp.VarOp, Deque<Object>> variableStack, CoreOp.VarOp vop) {
259         Object to = variableStack.get(vop).peek();
260         return throwIfUninitialized(vop, to);
261     }
262 
263     static Object throwIfUninitialized(CoreOp.VarOp vop, Object to) {
264         if (to instanceof Uninitialized) {
265             throw new IllegalStateException("Loading from uninitialized variable: " + vop);
266         }
267         return to;
268     }
269 
270     /**
271      * Finds the join points of a body.
272      * <p>
273      * A join point is a block that is in the dominance frontier of one or more predecessors, that make one or more
274      * stores to variables (using the {@code VarStoreOp} operation on the result of a {@code VarOp} operation).
275      * The join point contains the set variables ({@code VarOp} operations) that are stored to.
276      * <p>
277      * A variable of a joint point indicates that a block argument may need to be added to the join point's block
278      * when converting variables to SSA form. Different values of a variable may occur at different control flow
279      * paths at the join point. The block argument represents the convergence of multiple values for the same
280      * variable, where a predecessor assigns to the block argument.
281      * (Block arguments are equivalent to phi-values, or phi-nodes, used in other representations.)
282      *
283      * @param body the body.
284      * @param joinPoints the returned join points.
285      * @implNote See "Efficiently Computing Static Single Assignment Form and the Control Dependence Graph" by Ron Cytron et. al.
286      * Section 5.1 and Figure 11.
287      */
288     static void findJoinPoints(Body body, Map<Block, Set<CoreOp.VarOp>> joinPoints) {
289         Map<Block, Set<Block>> df = body.dominanceFrontier();
290         Map<CoreOp.VarOp, Set<Block>> a = findVarStores(body);
291 
292         int iterCount = 0;
293         int[] hasAlready = new int[body.blocks().size()];
294         int[] work = new int[body.blocks().size()];
295 
296         Deque<Block> w = new ArrayDeque<>();
297 
298         for (CoreOp.VarOp v : a.keySet()) {
299             iterCount++;
300 
301             for (Block x : a.get(v)) {
302                 work[x.index()] = iterCount;
303                 w.push(x);
304             }
305 
306             while (!w.isEmpty()) {
307                 Block x = w.pop();
308 
309                 for (Block y : df.getOrDefault(x, Set.of())) {
310                     if (hasAlready[y.index()] < iterCount) {
311                         // Only add to the join points if y is dominated by the var's block
312                         if (y.isDominatedBy(v.parentBlock())) {
313                             joinPoints.computeIfAbsent(y, _k -> new LinkedHashSet<>()).add(v);
314                         }
315                         hasAlready[y.index()] = iterCount;
316 
317                         if (work[y.index()] < iterCount) {
318                             work[y.index()] = iterCount;
319                             w.push(y);
320                         }
321                     }
322                 }
323             }
324         }
325     }
326 
327     // Returns map of variable to blocks that contain stores to the variables declared in the body
328     // Throws ISE if a descendant store operation is encountered
329     // @@@ Compute map for whole tree, then traverse keys with filter
330     static Map<CoreOp.VarOp, Set<Block>> findVarStores(Body r) {
331         return r.traverse(new LinkedHashMap<>(), CodeElement.opVisitor((stores, op) -> {
332             if (op instanceof CoreOp.VarAccessOp.VarStoreOp storeOp) {
333                 if (storeOp.varOp().ancestorBody() != storeOp.ancestorBody()) {
334                     throw new IllegalStateException("Descendant variable store operation");
335                 }
336                 if (storeOp.varOp().ancestorBody() == r) {
337                     stores.computeIfAbsent(storeOp.varOp(), _v -> new LinkedHashSet<>()).add(storeOp.parentBlock());
338                 }
339             }
340             return stores;
341         }));
342     }
343 
344     record Node(Block b, Set<Node> children) {
345     }
346 
347     static Node buildDomTree(Block entryBlock, Map<Block, Block> idoms) {
348         Map<Block, Node> tree = new HashMap<>();
349         for (Map.Entry<Block, Block> e : idoms.entrySet()) {
350             Block id = e.getValue();
351             Block b = e.getKey();
352 
353             Node parent = tree.computeIfAbsent(id, _k -> new Node(_k, new HashSet<>()));
354             if (b == entryBlock) {
355                 continue;
356             }
357 
358             Node child = tree.computeIfAbsent(b, _k -> new Node(_k, new HashSet<>()));
359             parent.children.add(child);
360         }
361         return tree.get(entryBlock);
362     }
363 }