Improve Bytecode DSL code units dumping

Author: steve-s

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/code/PCode.java

@@ -685,10 +685,8 @@ public String toString() {
     public String toDisassembledString(boolean quickened) {
         RootNode rootNode = getRootCallTarget().getRootNode();
         if (PythonOptions.ENABLE_BYTECODE_DSL_INTERPRETER && rootNode instanceof PBytecodeDSLRootNode dslRoot) {
-            return dslRoot.dump();
-        }
-
-        if (rootNode instanceof PBytecodeGeneratorRootNode r) {
+            return dslRoot.getCodeUnit().toString(quickened);
+        } else if (rootNode instanceof PBytecodeGeneratorRootNode r) {
             rootNode = r.getBytecodeRootNode();
         } else if (rootNode instanceof PBytecodeGeneratorFunctionRootNode r) {
             rootNode = r.getBytecodeRootNode();

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/compiler/BytecodeCodeUnit.java

@@ -125,34 +125,11 @@ public int bciToColumn(int bci) {
     }
 
     @Override
-    public String toString() {
-        return toString(false);
-    }
-
-    public String toString(boolean quickened) {
-        StringBuilder sb = new StringBuilder();
-
-        HashMap<Integer, String[]> lines = new HashMap<>();
-
-        sb.append("Disassembly of ").append(qualname).append(":\n");
-
-        List<String> flagNames = new ArrayList<>();
-        if (isGenerator()) {
-            flagNames.add("CO_GENERATOR");
-        }
-        if (isCoroutine()) {
-            flagNames.add("CO_COROUTINE");
-        }
-        if (isAsyncGenerator()) {
-            flagNames.add("CO_ASYNC_GENERATOR");
-        }
-        if (!flagNames.isEmpty()) {
-            sb.append("Flags: ").append(String.join(" | ", flagNames)).append("\n");
-        }
-
+    protected void dumpBytecode(StringBuilder sb, boolean quickened) {
         int bci = 0;
         int oparg = 0;
         SourceMap map = getSourceMap();
+        HashMap<Integer, String[]> lines = new HashMap<>();
         while (bci < code.length) {
             int bcBCI = bci;
             OpCodes opcode = OpCodes.fromOpCode(code[bci++]);
@@ -393,15 +370,6 @@ public String toString(boolean quickened) {
                 sb.append('\n');
             }
         }
-
-        for (Object c : constants) {
-            if (c instanceof CodeUnit) {
-                sb.append('\n');
-                sb.append(c);
-            }
-        }
-
-        return sb.toString();
     }
 
     private static String collectionKindToString(int oparg) {

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/compiler/CodeUnit.java

@@ -40,7 +40,9 @@
  */
 package com.oracle.graal.python.compiler;
 
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.List;
 
 import com.oracle.graal.python.builtins.objects.code.PCode;
 import com.oracle.graal.python.builtins.objects.function.Signature;
@@ -167,4 +169,46 @@ public final Signature computeSignature() {
                         parameterNames,
                         kwOnlyNames);
     }
+
+    @Override
+    public String toString() {
+        return toString(false);
+    }
+
+    protected abstract void dumpBytecode(StringBuilder sb, boolean optimized);
+
+    /**
+     * @param optimized Whether to print the initial state of the bytecode or current state, if
+     *            available, where some instructions may be transformed, e.g., quickened.
+     */
+    public String toString(boolean optimized) {
+        StringBuilder sb = new StringBuilder();
+
+        sb.append("Disassembly of ").append(qualname).append(":\n");
+
+        List<String> flagNames = new ArrayList<>();
+        if (isGenerator()) {
+            flagNames.add("CO_GENERATOR");
+        }
+        if (isCoroutine()) {
+            flagNames.add("CO_COROUTINE");
+        }
+        if (isAsyncGenerator()) {
+            flagNames.add("CO_ASYNC_GENERATOR");
+        }
+        if (!flagNames.isEmpty()) {
+            sb.append("Flags: ").append(String.join(" | ", flagNames)).append("\n");
+        }
+
+        dumpBytecode(sb, optimized);
+
+        for (Object c : constants) {
+            if (c instanceof CodeUnit cd) {
+                sb.append('\n');
+                sb.append(cd.toString(optimized));
+            }
+        }
+
+        return sb.toString();
+    }
 }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode_dsl/BytecodeDSLCodeUnit.java

@@ -129,4 +129,14 @@ public TruffleString getDocstring() {
         return null;
     }
 
+    @Override
+    protected void dumpBytecode(StringBuilder sb, boolean optimized) {
+        for (int i = 0; i < nodes.count(); i++) {
+            if (i != 0) {
+                sb.append('\n');
+            }
+            sb.append(nodes.getNode(i).dump());
+            sb.append('\n'); // dump does not print newline at the end
+        }
+    }
 }