1 /*
  2  * Copyright (c) 2024, 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.extern;
 27 
 28 import java.io.*;
 29 
 30 import jdk.incubator.code.*;
 31 import jdk.incubator.code.dialect.java.JavaType;
 32 import jdk.incubator.code.dialect.java.impl.JavaTypeUtils;
 33 
 34 import java.lang.reflect.Array;
 35 import java.nio.charset.StandardCharsets;
 36 import java.util.HashMap;
 37 import java.util.Map;
 38 import java.util.function.Consumer;
 39 import java.util.function.Function;
 40 
 41 /**
 42  * A writer of code models to the textual form.
 43  * <p>
 44  * A code model in textual form may be parsed back into the runtime form by parsing it.
 45  */
 46 public final class OpWriter {
 47 
 48     /**
 49      * The attribute name associated with the location attribute.
 50      */
 51     static final String ATTRIBUTE_LOCATION = "loc";
 52 
 53     static final class GlobalValueBlockNaming implements Function<CodeItem, String> {
 54         final Map<CodeItem, String> gn;
 55         int valueOrdinal = 0;
 56 
 57         GlobalValueBlockNaming() {
 58             this.gn = new HashMap<>();
 59         }
 60 
 61         private String name(Block b) {
 62             Block p = b.ancestorBlock();
 63             return (p == null ? "block_" : name(p) + "_") + b.index();
 64         }
 65 
 66         @Override
 67         public String apply(CodeItem codeItem) {
 68             return switch (codeItem) {
 69                 case Block block -> gn.computeIfAbsent(block, _b -> name(block));
 70                 case Value value -> gn.computeIfAbsent(value, _v -> String.valueOf(valueOrdinal++));
 71                 default -> throw new IllegalStateException("Unexpected code item: " + codeItem);
 72             };
 73         }
 74     }
 75 
 76     static final class AttributeMapper {
 77         static String toString(Object value) {
 78             if (value == ExternalizedOp.NULL_ATTRIBUTE_VALUE) {
 79                 return "null";
 80             }
 81 
 82             StringBuilder sb = new StringBuilder();
 83             toString(value, sb);
 84             return sb.toString();
 85         }
 86 
 87         static void toString(Object o, StringBuilder sb) {
 88             if (o.getClass().isArray()) {
 89                 // note, while we can't parse back the array representation, this might be useful
 90                 // for non-externalizable ops that want better string representation of array attribute values (e.g. ONNX)
 91                 arrayToString(o, sb);
 92             } else {
 93                 switch (o) {
 94                     case Integer i -> sb.append(i);
 95                     case Long l -> sb.append(l).append('L');
 96                     case Float f -> sb.append(f).append('f');
 97                     case Double d -> sb.append(d).append('d');
 98                     case Character c -> sb.append('\'').append(c).append('\'');
 99                     case Boolean b -> sb.append(b);
100                     case TypeElement te -> sb.append(JavaTypeUtils.flatten(te.externalize()));
101                     default -> {  // fallback to a string
102                         sb.append('"');
103                         quote(o.toString(), sb);
104                         sb.append('"');
105                     }
106                 }
107             }
108         }
109 
110         static void arrayToString(Object a, StringBuilder sb) {
111             boolean first = true;
112             sb.append("[");
113             for (int i = 0; i < Array.getLength(a); i++) {
114                 if (!first) {
115                     sb.append(", ");
116                 }
117 
118                 toString(Array.get(a, i), sb);
119                 first = false;
120             }
121             sb.append("]");
122         }
123     }
124 
125     static void quote(String s, StringBuilder sb) {
126         for (int i = 0; i < s.length(); i++) {
127             sb.append(quote(s.charAt(i)));
128         }
129     }
130 
131     /**
132      * Escapes a character if it has an escape sequence or is
133      * non-printable ASCII.  Leaves non-ASCII characters alone.
134      */
135     // Copied from com.sun.tools.javac.util.Convert
136     static String quote(char ch) {
137         return switch (ch) {
138             case '\b' -> "\\b";
139             case '\f' -> "\\f";
140             case '\n' -> "\\n";
141             case '\r' -> "\\r";
142             case '\t' -> "\\t";
143             case '\'' -> "\\'";
144             case '\"' -> "\\\"";
145             case '\\' -> "\\\\";
146             default -> (isPrintableAscii(ch))
147                     ? String.valueOf(ch)
148                     : String.format("\\u%04x", (int) ch);
149         };
150     }
151 
152     /**
153      * Is a character printable ASCII?
154      */
155     static boolean isPrintableAscii(char ch) {
156         return ch >= ' ' && ch <= '~';
157     }
158 
159     static final class IndentWriter extends Writer {
160         static final int INDENT = 2;
161 
162         final Writer w;
163         int indent;
164         boolean writeIndent = true;
165 
166         IndentWriter(Writer w) {
167             this(w, 0);
168         }
169 
170         IndentWriter(Writer w, int indent) {
171             this.w = w;
172             this.indent = indent;
173         }
174 
175         @Override
176         public void write(char[] cbuf, int off, int len) throws IOException {
177             if (writeIndent) {
178                 w.write(" ".repeat(indent));
179                 writeIndent = false;
180             }
181             w.write(cbuf, off, len);
182             if (len > 0 && cbuf[off + len - 1] == '\n') {
183                 writeIndent = true;
184             }
185         }
186 
187         @Override
188         public void flush() throws IOException {
189             w.flush();
190         }
191 
192         @Override
193         public void close() throws IOException {
194             w.close();
195         }
196 
197         void in() {
198             in(INDENT);
199         }
200 
201         void in(int i) {
202             indent += i;
203         }
204 
205         void out() {
206             out(INDENT);
207         }
208 
209         void out(int i) {
210             indent -= i;
211         }
212     }
213 
214     /**
215      * Computes global names for blocks and values in a code model.
216      * <p>
217      * The code model is traversed in the same order as if the model
218      * was written. Therefore, the names in the returned map will the
219      * same as the names that are written. This can be useful for debugging
220      * and testing.
221      *
222      * @param root the code model
223      * @return the map of computed names, modifiable
224      */
225     public static Function<CodeItem, String> computeGlobalNames(Op root) {
226         OpWriter w = new OpWriter(Writer.nullWriter());
227         w.writeOp(root);
228         return w.namer();
229     }
230 
231     /**
232      * Writes a code model (an operation) to the output stream, using the UTF-8 character set.
233      *
234      * @param out the output stream
235      * @param op the code model
236      */
237     public static void writeTo(OutputStream out, Op op, Option... options) {
238         writeTo(new OutputStreamWriter(out, StandardCharsets.UTF_8), op, options);
239     }
240 
241     /**
242      * Writes a code model (an operation) to the character stream.
243      * <p>
244      * The character stream will be flushed after the model is writen.
245      *
246      * @param w the character stream
247      * @param op the code model
248      * @param options the writer options
249      */
250     public static void writeTo(Writer w, Op op, Option... options) {
251         OpWriter ow = new OpWriter(w, options);
252         ow.writeOp(op);
253         try {
254             // @@@ Is this needed?
255             w.flush();
256         } catch (IOException e) {
257             throw new UncheckedIOException(e);
258         }
259     }
260 
261     /**
262      * Writes a code model (an operation) to a string.
263      *
264      * @param op the code model
265      * @param options the writer options
266      */
267     public static String toText(Op op, OpWriter.Option... options) {
268         StringWriter w = new StringWriter();
269         writeTo(w, op, options);
270         return w.toString();
271     }
272 
273     /**
274      * An option that affects the writing operations.
275      */
276     public sealed interface Option {
277     }
278 
279     /**
280      * An option describing the function to use for naming code items.
281      */
282     public sealed interface CodeItemNamerOption extends Option
283             permits NamerOptionImpl {
284 
285         static CodeItemNamerOption of(Function<CodeItem, String> named) {
286             return new NamerOptionImpl(named);
287         }
288 
289         static CodeItemNamerOption defaultValue() {
290             return of(new GlobalValueBlockNaming());
291         }
292 
293         Function<CodeItem, String> namer();
294     }
295     private record NamerOptionImpl(Function<CodeItem, String> namer) implements CodeItemNamerOption {
296     }
297 
298     /**
299      * An option describing whether location information should be written or dropped.
300      */
301     public enum LocationOption implements Option {
302         /** Writes location */
303         WRITE_LOCATION,
304         /** Drops location */
305         DROP_LOCATION;
306 
307         public static LocationOption defaultValue() {
308             return WRITE_LOCATION;
309         }
310     }
311 
312     /**
313      * An option describing whether an operation's descendant code elements should be written or dropped.
314      */
315     public enum OpDescendantsOption implements Option {
316         /** Writes descendants of an operation, if any */
317         WRITE_DESCENDANTS,
318         /** Drops descendants of an operation, if any */
319         DROP_DESCENDANTS;
320 
321         public static OpDescendantsOption defaultValue() {
322             return WRITE_DESCENDANTS;
323         }
324     }
325 
326     /**
327      * An option describing whether an operation's result be written or dropped if its type is void.
328      */
329     public enum VoidOpResultOption implements Option {
330         /** Writes void operation result */
331         WRITE_VOID,
332         /** Drops void operation result */
333         DROP_VOID;
334 
335         public static VoidOpResultOption defaultValue() {
336             return DROP_VOID;
337         }
338     }
339 
340     final Function<CodeItem, String> namer;
341     final IndentWriter w;
342     final boolean dropLocation;
343     final boolean dropOpDescendants;
344     final boolean writeVoidOpResult;
345 
346     /**
347      * Creates a writer of code models (operations) to their textual form.
348      *
349      * @param w the character stream writer to write the textual form.
350      */
351     public OpWriter(Writer w) {
352         this.w = new IndentWriter(w);
353         this.namer = new GlobalValueBlockNaming();
354         this.dropLocation = false;
355         this.dropOpDescendants = false;
356         this.writeVoidOpResult = false;
357     }
358 
359     /**
360      * Creates a writer of code models (operations) to their textual form.
361      *
362      * @param w the character stream writer to write the textual form.
363      * @param options the writer options
364      */
365     public OpWriter(Writer w, Option... options) {
366         Function<CodeItem, String> namer = null;
367         boolean dropLocation = false;
368         boolean dropOpDescendants = false;
369         boolean writeVoidOpResult = false;
370         for (Option option : options) {
371             switch (option) {
372                 case CodeItemNamerOption namerOption -> {
373                     namer = namerOption.namer();
374                 }
375                 case LocationOption locationOption -> {
376                     dropLocation = locationOption ==
377                             LocationOption.DROP_LOCATION;
378                 }
379                 case OpDescendantsOption opDescendantsOption -> {
380                     dropOpDescendants = opDescendantsOption ==
381                             OpDescendantsOption.DROP_DESCENDANTS;
382                 }
383                 case VoidOpResultOption voidOpResultOption -> {
384                     writeVoidOpResult = voidOpResultOption == VoidOpResultOption.WRITE_VOID;
385                 }
386             }
387         }
388 
389         this.w = new IndentWriter(w);
390         this.namer = (namer == null) ? new GlobalValueBlockNaming() : namer;
391         this.dropLocation = dropLocation;
392         this.dropOpDescendants = dropOpDescendants;
393         this.writeVoidOpResult = writeVoidOpResult;
394     }
395 
396     /**
397      * {@return the function that names blocks and values.}
398      */
399     public Function<CodeItem, String> namer() {
400         return namer;
401     }
402 
403     /**
404      * Writes a code model, an operation, to the character stream.
405      *
406      * @param op the code model
407      */
408     public void writeOp(Op op) {
409         if (op.parent() != null) {
410             Op.Result opr = op.result();
411             if (writeVoidOpResult || !opr.type().equals(JavaType.VOID)) {
412                 writeValueDeclaration(opr);
413                 write(" = ");
414             }
415         }
416         write(op.opName());
417 
418         if (!op.operands().isEmpty()) {
419             write(" ");
420             writeSpaceSeparatedList(op.operands(), this::writeValueUse);
421         }
422 
423         if (!op.successors().isEmpty()) {
424             write(" ");
425             writeSpaceSeparatedList(op.successors(), this::writeSuccessor);
426         }
427 
428         if (!dropLocation) {
429             Location location = op.location();
430             if (location != null) {
431                 write(" ");
432                 writeAttribute(ATTRIBUTE_LOCATION, op.location());
433             }
434         }
435         Map<String, Object> attributes = op.externalize();
436         if (!attributes.isEmpty()) {
437             write(" ");
438             writeSpaceSeparatedList(attributes.entrySet(), e -> writeAttribute(e.getKey(), e.getValue()));
439         }
440 
441         if (!dropOpDescendants && !op.bodies().isEmpty()) {
442             int nBodies = op.bodies().size();
443             if (nBodies == 1) {
444                 write(" ");
445             } else {
446                 write("\n");
447                 w.in();
448                 w.in();
449             }
450             boolean first = true;
451             for (Body body : op.bodies()) {
452                 if (!first) {
453                     write("\n");
454                 }
455                 writeBody(body);
456                 first = false;
457             }
458             if (nBodies > 1) {
459                 w.out();
460                 w.out();
461             }
462         }
463 
464         write(";");
465     }
466 
467     void writeSuccessor(Block.Reference successor) {
468         writeBlockName(successor.targetBlock());
469         if (!successor.arguments().isEmpty()) {
470             write("(");
471             writeCommaSeparatedList(successor.arguments(), this::writeValueUse);
472             write(")");
473         }
474     }
475 
476     void writeAttribute(String name, Object value) {
477         write("@");
478         if (!name.isEmpty()) {
479             write(name);
480             write("=");
481         }
482         write(AttributeMapper.toString(value));
483     }
484 
485     void writeBody(Body body) {
486         Block eb = body.entryBlock();
487         write("(");
488         writeCommaSeparatedList(eb.parameters(), this::writeValueDeclaration);
489         write(")");
490         writeType(body.bodyType().returnType());
491         write(" -> {\n");
492         w.in();
493         for (Block b : body.blocks()) {
494             if (!b.isEntryBlock()) {
495                 write("\n");
496             }
497             writeBlock(b);
498         }
499         w.out();
500         write("}");
501     }
502 
503     void writeBlock(Block block) {
504         if (!block.isEntryBlock()) {
505             writeBlockName(block);
506             if (!block.parameters().isEmpty()) {
507                 write("(");
508                 writeCommaSeparatedList(block.parameters(), this::writeValueDeclaration);
509                 write(")");
510             }
511             write(":\n");
512         }
513         w.in();
514         for (Op op : block.ops()) {
515             writeOp(op);
516             write("\n");
517         }
518         w.out();
519     }
520 
521     void writeBlockName(Block b) {
522         write("^");
523         write(namer.apply(b));
524     }
525 
526     void writeValueUse(Value v) {
527         write("%");
528         write(namer.apply(v));
529     }
530 
531     void writeValueDeclaration(Value v) {
532         write("%");
533         write(namer.apply(v));
534         write(" : ");
535         writeType(v.type());
536     }
537 
538     <T> void writeSpaceSeparatedList(Iterable<T> l, Consumer<T> c) {
539         writeSeparatedList(" ", l, c);
540     }
541 
542     <T> void writeCommaSeparatedList(Iterable<T> l, Consumer<T> c) {
543         writeSeparatedList(", ", l, c);
544     }
545 
546     <T> void writeSeparatedList(String separator, Iterable<T> l, Consumer<T> c) {
547         boolean first = true;
548         for (T t : l) {
549             if (!first) {
550                 write(separator);
551             }
552             c.accept(t);
553             first = false;
554         }
555     }
556 
557     void writeType(TypeElement te) {
558         write(JavaTypeUtils.flatten(te.externalize()).toString());
559     }
560 
561     void write(String s) {
562         try {
563             w.write(s);
564         } catch (IOException e) {
565             throw new UncheckedIOException(e);
566         }
567     }
568 }