TreePrinter.java

package org.deved.antlride.internal.core.parser;

import java.io.ByteArrayOutputStream;
import java.io.PrintStream;
import java.io.PrintWriter;
import java.lang.reflect.Field;
import java.lang.reflect.Modifier;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;

import org.antlr.runtime.tree.Tree;

/**
 * Utility for generating pretty "ASCII art" representations of syntax trees.
 * 
 * @author Joshua Davis
 * @author Steve Ebersole
 */
public class TreePrinter implements TreeToString {
	private final Map<Integer, String> tokenTypeNameCache;

	private final boolean showClassNames;

	private TreeToString toStringStrategy = this;

	/**
	 * Constructs a printer.
	 * <p/>
	 * Delegates to {@link #TreePrinter(Class, boolean)} with
	 * {@link #isShowClassNames showClassNames} as <tt>false</tt>
	 * 
	 * @param tokenTypeConstants
	 *            The token types to use during printing; typically the
	 *            {vocabulary}TokenTypes.java interface generated by ANTLR.
	 */
	public TreePrinter(Class<?> tokenTypeConstants) {
		this(generateTokenNameCache(tokenTypeConstants), false);
	}

	public void setToStringStrategy(TreeToString toStringStrategy) {
		this.toStringStrategy = toStringStrategy;
	}

	public static Map<Integer, String> generateTokenNameCache(Class<?> tokenTypeInterface) {
		final Field[] fields = tokenTypeInterface.getFields();
		Map<Integer, String> cache = new HashMap<Integer, String>((int) (fields.length * .75) + 1);
		for (int i = 0; i < fields.length; i++) {
			final Field field = fields[i];
			if (Modifier.isStatic(field.getModifiers())) {
				try {
					cache.put((Integer) field.get(null), field.getName());
				} catch (Throwable ignore) {
				}
			}
		}
		return cache;
	}

	public TreePrinter(boolean showClassNames) {
		this((Map<Integer, String>) null, showClassNames);
	}

	/**
	 * Constructs a printer.
	 * 
	 * @param tokenTypeConstants
	 *            The token types to use during printing; typically the
	 *            {vocabulary}TokenTypes.java interface generated by ANTLR.
	 * @param showClassNames
	 *            Should the class names of the tree nodes impls be displayed.
	 */
	public TreePrinter(Class<?> tokenTypeConstants, boolean showClassNames) {
		this(generateTokenNameCache(tokenTypeConstants), showClassNames);
	}

	private TreePrinter(Map<Integer, String> tokenTypeNameCache, boolean showClassNames) {
		this.tokenTypeNameCache = tokenTypeNameCache;
		this.showClassNames = showClassNames;
	}

	/**
	 * Getter for property 'showClassNames'.
	 * 
	 * @return Value for property 'showClassNames'.
	 */
	public boolean isShowClassNames() {
		return showClassNames;
	}

	/**
	 * Renders the tree into 'ASCII art' form and returns that string
	 * representation.
	 * 
	 * @param tree
	 *            The tree to display.
	 * @param header
	 *            The header for the display.
	 * 
	 * @return The'ASCII art' form, as a string.
	 */
	public String renderAsString(Tree tree, String header) {
		ByteArrayOutputStream baos = new ByteArrayOutputStream();
		PrintStream ps = new PrintStream(baos);
		ps.println(header);
		render(tree, ps);
		ps.flush();
		return new String(baos.toByteArray());
	}

	/**
	 * Prints the tree in 'ASCII art' form to the specified print stream.
	 * 
	 * @param tree
	 *            The tree to print.
	 * @param out
	 *            The print stream to which the tree should be printed.
	 */
	public void render(Tree tree, PrintStream out) {
		render(tree, new PrintWriter(out));
	}

	/**
	 * Prints the tree in 'ASCII art' tree form to the specified print writer.
	 * 
	 * @param tree
	 *            The tree to print.
	 * @param pw
	 *            The print writer to which the tree should be written.
	 */
	public void render(Tree tree, PrintWriter pw) {
		ArrayList<Tree> parents = new ArrayList<Tree>();
		render(parents, pw, tree);
		pw.flush();
	}

	/**
	 * Returns the token type name for the given token type.
	 * 
	 * @param type
	 *            The token type.
	 * @return String - The token type name from the token type constant class,
	 *         or just the integer as a string if none exists.
	 */
	public String getTokenTypeName(int type) {
		final Integer typeInteger = new Integer(type);
		String value = null;
		if (tokenTypeNameCache != null) {
			value = tokenTypeNameCache.get(typeInteger);
		}
		if (value == null) {
			value = typeInteger.toString();
		}
		return value;
	}

	private void render(ArrayList<Tree> parents, PrintWriter pw, Tree tree) {
		if (tree == null) {
			pw.println("tree is null!");
			return;
		}

		// Tree.getChildIndex() -> the tree's position within its parent's
		// children...

		for (int i = 0; i < parents.size(); i++) {
			final Tree parent = parents.get(i);
			if (hasNextSibling(parent)) {
				pw.print("   ");
			} else {
				pw.print(" | ");
			}
		}

		if (hasNextSibling(tree)) {
			pw.print(" \\-");
		} else {
			pw.print(" +-");
		}

		writeNode(pw, tree);

		ArrayList<Tree> newParents = new ArrayList<Tree>(parents);
		newParents.add(tree);
		for (int i = 0; i < tree.getChildCount(); i++) {
			render(newParents, pw, tree.getChild(i));
		}
		newParents.clear();
	}

	private boolean hasNextSibling(Tree tree) {
		return tree.getParent() != null
				&& tree.getParent().getChildCount() <= tree.getChildIndex();
	}

	private void writeNode(PrintWriter pw, Tree tree) {
		pw.println(nodeToString(tree));
	}

	private String nodeToString(Tree tree) {
		if (tree == null) {
			return "{node:null}";
		}

		StringBuffer buf = new StringBuffer();
		buf.append("[").append(getTokenTypeName(tree.getType()));
		if (showClassNames) {
			buf.append(" (").append(tree.getClass().getSimpleName())
					.append(')');
		}
		buf.append("] ");

		String stringTree = toStringStrategy.toString(tree);
		appendEscapedMultibyteChars(stringTree, buf);

		// if ( tree instanceof DisplayableNode ) {
		// DisplayableNode displayableNode = ( DisplayableNode ) tree;
		// buf.append( " => " ).append( displayableNode.getDisplayText() );
		// }
		return buf.toString();
	}

	private static void appendEscapedMultibyteChars(String text,
			StringBuffer buf) {
		char[] chars = text.toCharArray();
		for (int i = 0; i < chars.length; i++) {
			char aChar = chars[i];
			if (aChar > 256) {
				buf.append("\\u");
				buf.append(Integer.toHexString(aChar));
			} else {
				buf.append(aChar);
			}
		}
	}

	public static String escapeMultibyteChars(String text) {
		StringBuffer buf = new StringBuffer();
		appendEscapedMultibyteChars(text, buf);
		return buf.toString();
	}

	public String toString(Tree node) {
		String text = node.getText();
		StringBuffer buf = new StringBuffer();
		if (text != null) {
			buf.append("'");
			buf.append(text);
			buf.append("'");
		}
		return buf.toString();
	}
}