Decompiler: Add code folding

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

@@ -49,9 +49,23 @@ public class ClangToken implements ClangNode {
 	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/component/ClangLayoutController.java

@@ -186,8 +186,10 @@ 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) {
 				AttributedString prototype = new AttributedString("prototype", color, metrics);
 				Program program = decompilerPanel.getProgram();
@@ -226,8 +228,8 @@ private Color getFunctionColor(Function function) {
 
 		// For now we have decided that any external function, linked or not, will be one color, as
 		// this makes it easy for the user to identify external function calls. Other functions will
-		// be colored according to the SymbolInspector.  If we use the SymbolInspector for all 
-		// colors, then some of the color values will be very close to some of the colors used by 
+		// be colored according to the SymbolInspector.  If we use the SymbolInspector for all
+		// colors, then some of the color values will be very close to some of the colors used by
 		// the Decompiler.  For example, non-linked external functions default to red and linked
 		// external functions default to green.
 		if (function.isExternal()) {

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

@@ -55,6 +55,7 @@
 import ghidra.util.bean.field.AnnotatedTextFieldElement;
 import ghidra.util.task.SwingUpdateManager;
 
+
 /**
  * Class to handle the display of a decompiled function
  */
@@ -158,7 +159,7 @@ public void componentResized(ComponentEvent e) {
 		setDecompileData(new EmptyDecompileData("No Function"));
 
 		if (options.isDisplayLineNumbers()) {
-			addMarginProvider(lineNumbersMargin = new LineNumberDecompilerMarginProvider());
+			addMarginProvider(lineNumbersMargin = new LineNumberDecompilerMarginProvider(this));
 		}
 	}
 
@@ -786,6 +787,94 @@ else if (DockingUtils.isControlModifier(ev) && ev.isShiftDown()) {
 		}
 	}
 
+	public void arrowClickAction(int y) {
+		int lineNumber = getLineNumber(y);
+		ClangToken openingBraceToken = null;
+		ClangLine line = getLines().get(lineNumber - 1);
+		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;
+			for (ClangNode element : list) {
+				ClangToken token = (ClangToken) element;
+				if (inSection) {
+					if ((token instanceof ClangSyntaxToken)) {
+						inSection = (!token.equals(closingBrace));
+					}
+					if (inSection) {
+						token.setCollapsedToken(!isCollapsed);
+					}
+				}
+				else if ((token instanceof ClangSyntaxToken)) {
+					inSection = (token.equals(openingBrace));
+				}
+			}
+
+			setDecompileData(decompileData);
+		}
+	}
+
+	private boolean isBlockCollapsed(ClangSyntaxToken openingBrace) {
+		ClangSyntaxToken closingBrace = DecompilerUtils.getMatchingBrace(openingBrace);
+		if (closingBrace == null) {
+			return false;
+		}
+
+		List<ClangNode> list = new ArrayList<>();
+		openingBrace.Parent().flatten(list);
+
+		boolean inSection = false;
+		for (ClangNode element : list) {
+			ClangToken token = (ClangToken) element;
+			if (inSection) {
+				if (token.equals(closingBrace)) {
+					break;
+				}
+				return token.getCollapsedToken();
+			} else if (token.equals(openingBrace)) {
+				inSection = true;
+			}
+		}
+		return false;
+	}
+
+	public Map<Integer, Boolean> getLinesWithOpeningBraces() {
+		Map<Integer, Boolean> lineNumbers = 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);
+					lineNumbers.put(i, isBlockCollapsed((ClangSyntaxToken) token));
+				}
+			}
+		}
+		return lineNumbers;
+	}
+
 	private void tryToGoto(FieldLocation location, Field field, MouseEvent event,
 			boolean newWindow) {
 		if (!navigationEnabled) {
@@ -1308,7 +1397,7 @@ public void optionsChanged(DecompileOptions decompilerOptions) {
 
 		if (options.isDisplayLineNumbers()) {
 			if (lineNumbersMargin == null) {
-				addMarginProvider(lineNumbersMargin = new LineNumberDecompilerMarginProvider());
+				addMarginProvider(lineNumbersMargin = new LineNumberDecompilerMarginProvider(this));
 			}
 		}
 		else {

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;

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

@@ -16,7 +16,10 @@
 package ghidra.app.decompiler.component.margin;
 
 import java.awt.*;
+import java.awt.event.MouseAdapter;
+import java.awt.event.MouseEvent;
 import java.math.BigInteger;
+import java.util.Map;
 
 import javax.swing.BorderFactory;
 import javax.swing.JPanel;
@@ -26,6 +29,7 @@
 import docking.widgets.fieldpanel.listener.IndexMapper;
 import docking.widgets.fieldpanel.listener.LayoutModelListener;
 import ghidra.app.decompiler.DecompileOptions;
+import ghidra.app.decompiler.component.DecompilerPanel;
 import ghidra.program.model.listing.Program;
 
 /**
@@ -36,9 +40,17 @@ public class LineNumberDecompilerMarginProvider extends JPanel
 
 	private LayoutPixelIndexMap pixmap;
 	private LayoutModel model;
+	private final DecompilerPanel decompilerPanel;
 
-	public LineNumberDecompilerMarginProvider() {
+	public LineNumberDecompilerMarginProvider(DecompilerPanel decompilerPanel) {
+		this.decompilerPanel = decompilerPanel;
 		setBorder(BorderFactory.createEmptyBorder(0, 0, 0, 2));
+		addMouseListener(new MouseAdapter() {
+			@Override
+			public void mouseClicked(MouseEvent e) {
+				handleMouseClick(e);
+			}
+		});
 	}
 
 	@Override
@@ -91,26 +103,50 @@ private void setWidthForLastLine() {
 		}
 		int lastLine = model.getNumIndexes().intValueExact();
 		int width = getFontMetrics(getFont()).stringWidth(Integer.toString(lastLine));
+		int widthForArrows = getFontMetrics(getFont()).stringWidth(" ") * 2;
+		width += widthForArrows;
 		Insets insets = getInsets();
 		width += insets.left + insets.right;
-		setPreferredSize(new Dimension(Math.max(16, width), 0));
+		setPreferredSize(new Dimension(Math.max(32, width), 0));
 		invalidate();
 	}
 
+	private void handleMouseClick(MouseEvent e) {
+		Insets insets = getInsets();
+		int y = e.getY() - insets.top;
+		int x = e.getX() - insets.left;
+
+		if (x >= getWidth() - getFontMetrics(getFont()).stringWidth(" ") * 2 - insets.right) {
+			decompilerPanel.arrowClickAction(y);
+			repaint();
+		}
+	}
+
 	@Override
 	public void paint(Graphics g) {
 		super.paint(g);
 
 		Insets insets = getInsets();
-		int rightEdge = getWidth() - insets.right;
+		int rightEdge = getWidth() - insets.right - getFontMetrics(getFont()).stringWidth(" ");
+		int leftEdge = insets.left;
 		Rectangle visible = getVisibleRect();
 		BigInteger startIdx = pixmap.getIndex(visible.y);
 		BigInteger endIdx = pixmap.getIndex(visible.y + visible.height);
 		int ascent = g.getFontMetrics().getMaxAscent();
+		Map<Integer, Boolean> linesIndexes = decompilerPanel.getLinesWithOpeningBraces();
+
+
 		for (BigInteger i = startIdx; i.compareTo(endIdx) <= 0; i = i.add(BigInteger.ONE)) {
 			String text = i.add(BigInteger.ONE).toString();
-			int width = g.getFontMetrics().stringWidth(text);
-			GraphicsUtils.drawString(this, g, text, rightEdge - width, pixmap.getPixel(i) + ascent);
+			GraphicsUtils.drawString(this, g, text, leftEdge, pixmap.getPixel(i) + ascent);
+
+			if (linesIndexes.containsKey(i.intValue())) {
+				if (linesIndexes.get(i.intValue())) {
+					GraphicsUtils.drawString(this, g, "+", rightEdge, pixmap.getPixel(i) + ascent);
+				} else {
+					GraphicsUtils.drawString(this, g, "-", rightEdge, pixmap.getPixel(i) + ascent);
+				}
+			}
 		}
 	}
 }