Decompiler: Add code folding

Ghidra/Features/Decompiler/src/main/java/ghidra/app/decompiler/ClangToken.java

@@ -45,13 +45,29 @@ public class ClangToken implements ClangNode {
 	public final static int SPECIAL_COLOR = 10;
 	public final static int MAX_COLOR = 11;
 
+	public final static String ELLIPSIS_TEXT = "…";
+
 	private ClangNode parent;
 	private ClangLine lineparent;
 	private String text;
 	private int syntax_type;
+	private int collapseLevel = 0;
 	private Color highlight; // Color to highlight with or null if no highlight
 	private boolean matchingToken;
 
+	public boolean getCollapsedToken() {
+		return collapseLevel > 0;
+	}
+
+	public void setCollapsedToken(boolean collapsedToken) {
+		if (collapsedToken) {
+			collapseLevel++;
+		}
+		else {
+			collapseLevel--;
+		}
+	}
+
 	public ClangToken(ClangNode par) {
 		parent = par;
 		text = null;

Ghidra/Features/Decompiler/src/main/java/ghidra/app/decompiler/ClangTokenGroup.java

@@ -151,6 +151,18 @@ else if (elem == ELEM_BLOCK.id()) {
 			else {
 				ClangToken tok = ClangToken.buildToken(elem, this, decoder, pfactory);
 				AddTokenGroup(tok);
+
+				// Add an ellipsis syntax token after open curly braces that is
+				// collapsed by default. It will be uncollapsed when the rest of
+				// the block is collapsed, indicating that tokens have been
+				// collapsed at that point.
+				if ((tok instanceof ClangSyntaxToken) && "{".equals(tok.getText())) {
+					ClangSyntaxToken ellipsis = new ClangSyntaxToken(
+						this, ClangToken.ELLIPSIS_TEXT, ClangToken.COMMENT_COLOR
+					);
+					ellipsis.setCollapsedToken(true);
+					AddTokenGroup(ellipsis);
+				}
 			}
 			decoder.closeElement(elem);
 		}

Ghidra/Features/Decompiler/src/main/java/ghidra/app/decompiler/component/ClangLayoutController.java

@@ -186,6 +186,9 @@ private FieldElement[] createFieldElementsForLine(List<ClangToken> tokens) {
 		int columnPosition = 0;
 		for (int i = 0; i < tokens.size(); ++i) {
 			ClangToken token = tokens.get(i);
+			if (token.getCollapsedToken()) {
+				continue;
+			}
 			Color color = getTokenColor(token);
 
 			if (token instanceof ClangCommentToken) {

Ghidra/Features/Decompiler/src/main/java/ghidra/app/decompiler/component/DecompilerPanel.java

@@ -107,6 +107,15 @@ public class DecompilerPanel extends JPanel implements FieldMouseListener, Field
 
 	private DecompilerHoverProvider decompilerHoverProvider;
 
+	public class CodeBlock {
+		public int startLineIdx;  //< 0-based
+		public int numLines;
+		public ClangSyntaxToken openToken;
+	}
+
+	private Map<Integer, CodeBlock> blocks; // start line idx: num lines
+	private boolean pendingOptionChange = true;
+
 	DecompilerPanel(DecompilerController controller, DecompileOptions options,
 			DecompilerClipboardProvider clipboard, JComponent taskMonitorComponent) {
 		this.controller = controller;
@@ -158,7 +167,7 @@ public void componentResized(ComponentEvent e) {
 		setDecompileData(new EmptyDecompileData("No Function"));
 
 		if (options.isDisplayLineNumbers()) {
-			addMarginProvider(lineNumbersMargin = new LineNumberDecompilerMarginProvider());
+			addMarginProvider(lineNumbersMargin = new LineNumberDecompilerMarginProvider(this));
 		}
 	}
 
@@ -540,6 +549,13 @@ void setDecompileData(DecompileData decompileData) {
 		if (function != null) {
 			highlightController.reapplyAllHighlights(function);
 		}
+
+		// Only update the blocks when we're moving to a new function, or if the
+		// the display options changed
+		if (pendingOptionChange || !SystemUtilities.isEqual(oldData.getFunction(), decompileData.getFunction())) {
+			setBlocks();
+			pendingOptionChange = false;
+		}
 	}
 
 	private void setLocation(DecompileData oldData, DecompileData newData) {
@@ -786,6 +802,209 @@ else if (DockingUtils.isControlModifier(ev) && ev.isShiftDown()) {
 		}
 	}
 
+	public void arrowClickAction(int y) {
+		int linesIdx = pixmap.getIndex(y).intValue();
+		ClangToken openingBraceToken = null;
+		ClangLine line = getLines().get(linesIdx);
+		for (ClangToken lineToken : line.getAllTokens()) {
+			if ("{".equals(lineToken.getText())) {
+				openingBraceToken = lineToken;
+				break;
+			}
+		}
+
+		if (openingBraceToken instanceof ClangSyntaxToken) {
+			toggleCollapseToken((ClangSyntaxToken) openingBraceToken);
+		}
+	}
+
+	private void toggleCollapseToken(ClangSyntaxToken openingBrace) {
+		if (DecompilerUtils.isBrace(openingBrace)) {
+			ClangSyntaxToken closingBrace = DecompilerUtils.getMatchingBrace(openingBrace);
+			if (closingBrace == null) {
+				return;
+			}
+
+			boolean isCollapsed = isBlockCollapsed(openingBrace);
+			List<ClangNode> list = new ArrayList<>();
+			openingBrace.Parent().flatten(list);
+
+			boolean inSection = false;
+			boolean seenEllipsis = false;
+
+			for (ClangNode element : list) {
+				ClangToken token = (ClangToken) element;
+				if (token.equals(openingBrace)) {
+					inSection = true;
+					continue;
+				}
+
+				if (token.equals(closingBrace)) {
+					inSection = false;
+					break;
+				}
+
+				if (! inSection) {
+					continue;
+				}
+
+				boolean isEllipsis = (token instanceof ClangSyntaxToken) && ClangToken.ELLIPSIS_TEXT.equals(token.getText());
+
+				if (isEllipsis && !seenEllipsis) {
+					token.setCollapsedToken(isCollapsed);
+					seenEllipsis = true;
+				} else {
+					token.setCollapsedToken(!isCollapsed);
+				}
+			}
+
+			ClangToken cursorToken = getTokenAtCursor();
+			FieldLocation cursorPos = getCursorPosition();
+			boolean wasOnScreen = getOffscreenDistance(
+				cursorToken.getLineParent().getLineNumber() - 1
+			) == 0;
+
+			setDecompileData(decompileData);
+
+			if (! wasOnScreen) {
+				// If the cursor was not on screen to begin with, don't bother
+				// scrolling the cursor back into view.
+				return;
+			}
+
+			// Adjust the cursor to be at the same token
+			if (!cursorToken.getCollapsedToken()) {
+				// The token is still visible, it only moved to a different line
+				// The token only moved vertically, not horizontally, so we only
+				// adjust the row, not the column.
+				int lineNumber = cursorToken.getLineParent().getLineNumber() - 1;
+
+				fieldPanel.navigateTo(lineNumber, cursorPos.getCol());
+			} else {
+				// The token has been collapsed, iterate backwards until we find
+				// the first token that is not collapsed and target that
+
+				ClangNode cursorParent = cursorToken.Parent();
+				List<ClangNode> children = new ArrayList<>();
+				cursorParent.flatten(children);
+
+				int childIdx = children.indexOf(cursorToken);
+
+				// Attempt going up in the hierarchy at most 10 times
+				int j;
+				for (j = 0; j < 10; j++) {
+					int i;
+					for (i = childIdx; i >= 0; i--) {
+						ClangNode n = children.get(i);
+						if ((!(n instanceof ClangToken)) || ((ClangToken)n).getCollapsedToken()) {
+							continue;
+						}
+
+						cursorToken = (ClangToken)n;
+						break;
+					}
+
+					if (i >= 0) {
+						break;
+					}
+
+					// We did not find an uncollapsed token among these children,
+					// try the children of our parent.
+					ClangNode lastChild = children.get(0);
+					children.clear();
+					cursorParent = cursorParent.Parent();
+					cursorParent.flatten(children);
+					childIdx = children.indexOf(lastChild);
+				}
+
+				// Only change the token at the cursor if we managed to find an
+				// uncollapsed token. Otherwise, just leave the cursor at the
+				// same row/column position.
+				if (j != 10) {
+					goToToken(cursorToken);
+				}
+			}
+		}
+	}
+
+	public boolean isBlockCollapsed(ClangSyntaxToken openingBrace) {
+		ClangSyntaxToken closingBrace = DecompilerUtils.getMatchingBrace(openingBrace);
+		if (closingBrace == null) {
+			return false;
+		}
+
+		List<ClangNode> list = new ArrayList<>();
+		openingBrace.Parent().flatten(list);
+
+		// Check if the block is collapsed by checking the first token inside the
+		// block that is not an ellipsis syntax token.
+		boolean inSection = false;
+		for (ClangNode element : list) {
+			ClangToken token = (ClangToken) element;
+			if (!inSection) {
+				inSection = token.equals(openingBrace);
+				continue;
+			}
+
+			if (token.equals(closingBrace)) {
+				return false;
+			}
+
+			if ((token instanceof ClangSyntaxToken) && (ClangToken.ELLIPSIS_TEXT.equals(token.getText()))) {
+				continue;
+			}
+
+			return token.getCollapsedToken();
+		}
+
+		return false;
+	}
+
+	public Map<Integer, CodeBlock> getBlocks() {
+		return blocks;
+	}
+
+	// Note: this assumes there is only at most one block starting at any given address
+	public void setBlocks() {
+		blocks = new HashMap<>();
+		List<ClangLine> lines = getLines();
+
+		for (int i = 0; i < lines.size(); i++) {
+			List<ClangToken> lineTokens = lines.get(i).getAllTokens();
+
+			for (ClangToken token : lineTokens) {
+				if (token.getText().contains("{") && token instanceof ClangSyntaxToken) {
+					List<ClangNode> list = new ArrayList<>();
+					token.Parent().flatten(list);
+
+					// Determine opening and closing line numbers
+					ClangSyntaxToken open = (ClangSyntaxToken) token;
+					ClangSyntaxToken close = DecompilerUtils.getMatchingBrace(open);
+					Integer openLine = i;
+
+					// It must be possible to compute this more efficiently than
+					// O(n^2)...
+					Integer blockLen = 1;
+					for (int j = openLine; j < lines.size(); j++) {
+						if (lines.get(j).indexOfToken(close) != -1) {
+							// +1 because we want to do openLine + blockLen to
+							// get the line number after the closing brace
+							blockLen = j - openLine + 1;
+							break;
+						}
+					}
+
+					CodeBlock block = new CodeBlock();
+					block.startLineIdx = openLine;
+					block.numLines = blockLen;
+					block.openToken = open;
+
+					blocks.put(openLine, block);
+				}
+			}
+		}
+	}
+
 	private void tryToGoto(FieldLocation location, Field field, MouseEvent event,
 			boolean newWindow) {
 		if (!navigationEnabled) {
@@ -1225,7 +1444,20 @@ public ClangToken getTokenAtCursor() {
 	 * @return the line number, or 0 if not applicable
 	 */
 	public int getLineNumber(int y) {
-		return pixmap.getIndex(y).intValue() + 1;
+		int lineNumber = 0;
+		int idx = pixmap.getIndex(y).intValue();
+
+		for (int i = 0; i < idx; i++) {
+			CodeBlock block = blocks.getOrDefault(lineNumber, null);
+
+			if (block == null || !isBlockCollapsed(block.openToken)) {
+				lineNumber++;
+			} else {
+				lineNumber += block.numLines;
+			}
+		}
+
+		return lineNumber;
 	}
 
 	public DecompileOptions getOptions() {
@@ -1308,7 +1540,7 @@ public void optionsChanged(DecompileOptions decompilerOptions) {
 
 		if (options.isDisplayLineNumbers()) {
 			if (lineNumbersMargin == null) {
-				addMarginProvider(lineNumbersMargin = new LineNumberDecompilerMarginProvider());
+				addMarginProvider(lineNumbersMargin = new LineNumberDecompilerMarginProvider(this));
 			}
 		}
 		else {
@@ -1321,6 +1553,8 @@ public void optionsChanged(DecompileOptions decompilerOptions) {
 		for (DecompilerMarginProvider element : marginProviders) {
 			element.setOptions(options);
 		}
+
+		pendingOptionChange = true;
 	}
 
 	public void addMarginProvider(DecompilerMarginProvider provider) {

Ghidra/Features/Decompiler/src/main/java/ghidra/app/decompiler/component/DecompilerUtils.java

@@ -820,6 +820,9 @@ public static List<ClangLine> toLines(ClangTokenGroup group) {
 		for (; i < alltoks.size(); ++i) {
 
 			ClangToken tok = (ClangToken) alltoks.get(i);
+			if (tok.getCollapsedToken()) {
+				continue;
+			}
 			if (tok instanceof ClangBreak) {
 				lines.add(current);
 				brk = (ClangBreak) tok;