diff --git a/runtime/Java/src/org/antlr/v4/runtime/tree/gui/TreeViewer.java b/runtime/Java/src/org/antlr/v4/runtime/tree/gui/TreeViewer.java
index 28532f04d..dfd834fa1 100644
--- a/runtime/Java/src/org/antlr/v4/runtime/tree/gui/TreeViewer.java
+++ b/runtime/Java/src/org/antlr/v4/runtime/tree/gui/TreeViewer.java
@@ -299,37 +299,7 @@ public class TreeViewer extends JComponent {
 
 		JPanel wrapper = new JPanel(new FlowLayout());
 
-		// Add an export-to-png button left of the "OK" button
-		JButton png = new JButton("png");
-		png.addActionListener(
-			new ActionListener() {
-				@Override
-				public void actionPerformed(ActionEvent e) {
-					BufferedImage bi = new BufferedImage(viewer.getSize().width,
-							viewer.getSize().height, BufferedImage.TYPE_INT_ARGB);
-					Graphics g = bi.createGraphics();
-					viewer.paint(g);
-					g.dispose();
-
-					try{
-						File pngFile = new File("antlrv4_parse_tree_" +
-								System.currentTimeMillis() +  ".png");
-						ImageIO.write(bi, "png", pngFile);
-						JOptionPane.showMessageDialog(dialog,
-								"Saved PNG to: " + pngFile.getAbsolutePath());
-					} catch (Exception ex) {
-						JOptionPane.showMessageDialog(dialog,
-								"Could not export to PNG: " + ex.getMessage(),
-								"Error",
-								JOptionPane.ERROR_MESSAGE);
-						ex.printStackTrace();
-					}
-				}
-			}
-		);
-		wrapper.add(png);
-
-	  	// Add button to bottom
+		// Add button to bottom
 		JPanel bottomPanel = new JPanel(new BorderLayout(0,0));
 		contentPane.add(bottomPanel, BorderLayout.SOUTH);
 
@@ -345,6 +315,18 @@ public class TreeViewer extends JComponent {
 		);
 		wrapper.add(ok);
 
+		// Add an export-to-png button right of the "OK" button
+		JButton png = new JButton("png");
+		png.addActionListener(
+			new ActionListener() {
+				@Override
+				public void actionPerformed(ActionEvent e) {
+					generatePNGFile(viewer, dialog);
+				}
+			}
+		);
+		wrapper.add(png);
+
 		bottomPanel.add(wrapper, BorderLayout.SOUTH);
 
 		// Add scale slider
@@ -370,6 +352,30 @@ public class TreeViewer extends JComponent {
 		return dialog;
 	}
 
+	private static void generatePNGFile(TreeViewer viewer, JDialog dialog) {
+		BufferedImage bi = new BufferedImage(viewer.getSize().width,
+											 viewer.getSize().height,
+											 BufferedImage.TYPE_INT_ARGB);
+		Graphics g = bi.createGraphics();
+		viewer.paint(g);
+		g.dispose();
+
+		try{
+			File pngFile = new File("antlrv4_parse_tree_" +
+									System.currentTimeMillis() + ".png");
+			ImageIO.write(bi, "png", pngFile);
+			JOptionPane.showMessageDialog(dialog,
+										  "Saved PNG to: " + pngFile.getAbsolutePath());
+		}
+		catch (Exception ex) {
+			JOptionPane.showMessageDialog(dialog,
+										  "Could not export to PNG: " + ex.getMessage(),
+										  "Error",
+										  JOptionPane.ERROR_MESSAGE);
+			ex.printStackTrace();
+		}
+	}
+
 	private Dimension getScaledTreeSize() {
 		Dimension scaledTreeSize =
 			treeLayout.getBounds().getBounds().getSize();