Skip to content

Commit

Permalink
Class level assembler support for stack/variable analysis
Browse files Browse the repository at this point in the history
  • Loading branch information
@Col-E
Col-E committed Dec 13, 2023
1 parent 3d4f810 commit e67f619
Show file tree
Hide file tree
Showing 3 changed files with 165 additions and 49 deletions.
96 changes: 87 additions & 9 deletions recaf-ui/src/main/java/software/coley/recaf/ui/pane/editing/assembler/AssemblerPane.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,17 +2,22 @@

import jakarta.annotation.Nonnull;
import jakarta.enterprise.context.Dependent;
import jakarta.enterprise.inject.Instance;
import jakarta.inject.Inject;
import javafx.geometry.Orientation;
import javafx.scene.control.SplitPane;
import me.darknet.assembler.ast.ASTElement;
import me.darknet.assembler.ast.specific.ASTClass;
import me.darknet.assembler.ast.specific.ASTField;
import me.darknet.assembler.ast.specific.ASTMethod;
import me.darknet.assembler.compile.JavaClassRepresentation;
import me.darknet.assembler.compile.analysis.AnalysisException;
import me.darknet.assembler.compiler.ClassRepresentation;
import me.darknet.assembler.error.Error;
import me.darknet.assembler.error.Result;
import me.darknet.assembler.parser.Token;
import me.darknet.assembler.util.Location;
import org.fxmisc.richtext.CodeArea;
import org.slf4j.Logger;
import software.coley.recaf.analytics.logging.Logging;
import software.coley.recaf.info.ClassInfo;
Expand All @@ -25,6 +30,8 @@
import software.coley.recaf.services.navigation.UpdatableNavigable;
import software.coley.recaf.ui.LanguageStylesheets;
import software.coley.recaf.ui.config.KeybindingConfig;
import software.coley.recaf.ui.control.BoundTab;
import software.coley.recaf.ui.control.IconView;
import software.coley.recaf.ui.control.richtext.Editor;
import software.coley.recaf.ui.control.richtext.bracket.BracketMatchGraphicFactory;
import software.coley.recaf.ui.control.richtext.bracket.SelectedBracketTracking;
Expand All @@ -33,9 +40,8 @@
import software.coley.recaf.ui.control.richtext.syntax.RegexLanguages;
import software.coley.recaf.ui.control.richtext.syntax.RegexSyntaxHighlighter;
import software.coley.recaf.ui.pane.editing.AbstractContentPane;
import software.coley.recaf.util.Animations;
import software.coley.recaf.util.FxThreadUtil;
import software.coley.recaf.util.Unchecked;
import software.coley.recaf.ui.pane.editing.tabs.FieldsAndMethodsPane;
import software.coley.recaf.util.*;
import software.coley.recaf.workspace.model.bundle.Bundle;

import java.time.Duration;
Expand All @@ -61,6 +67,7 @@ public class AssemblerPane extends AbstractContentPane<PathNode<?>> implements U
private final ProblemTracking problemTracking = new ProblemTracking();
private final Editor editor = new Editor();
private final AtomicBoolean updateLock = new AtomicBoolean();
private final Instance<FieldsAndMethodsPane> fieldsAndMethodsPaneProvider;
private AssemblerPipeline<? extends ClassInfo, ? extends ClassRepresentation> pipeline;
private ClassRepresentation lastAssembledClassRepresentation;
private ClassInfo lastAssembledClass;
Expand All @@ -73,9 +80,11 @@ public class AssemblerPane extends AbstractContentPane<PathNode<?>> implements U
public AssemblerPane(@Nonnull AssemblerPipelineManager pipelineManager,
@Nonnull AssemblerToolTabs assemblerToolTabs,
@Nonnull SearchBar searchBar,
@Nonnull KeybindingConfig keys) {
@Nonnull KeybindingConfig keys,
@Nonnull Instance<FieldsAndMethodsPane> fieldsAndMethodsPaneProvider) {
this.pipelineManager = pipelineManager;
this.assemblerToolTabs = assemblerToolTabs;
this.fieldsAndMethodsPaneProvider = fieldsAndMethodsPaneProvider;

int timeToWait = pipelineManager.getServiceConfig().getDisassemblyAstParseDelay().getValue();

Expand All @@ -97,9 +106,6 @@ public AssemblerPane(@Nonnull AssemblerPipelineManager pipelineManager,
if (keys.getSave().match(event))
assembleAndUpdateWorkspace();
});

// TODO: For class level assemblers we need to track the 'editor' caret position and update the tool tabs
// to hold a ClassMemberPathNode to the current member the caret position is over
}

@Override
Expand Down Expand Up @@ -132,14 +138,86 @@ public ClassPathNode getClassPath() {

@Override
public void requestFocus(@Nonnull ClassMember member) {
if (lastRoughAst != null) {
System.out.println(lastRoughAst);
if (lastConcreteAst != null) {
for (ASTElement root : lastConcreteAst) {
if (root instanceof ASTClass astClass) {
for (ASTElement child : astClass.children()) {
String name;
String desc;
if (member.isMethod() && child instanceof ASTMethod astMethod) {
name = astMethod.getName().literal();
desc = astMethod.getDescriptor().literal();
} else if (member.isField() && child instanceof ASTField astField) {
name = astField.getName().literal();
desc = astField.getDescriptor().literal();
} else {
name = desc = null;
}
if (member.getName().equals(name) && member.getDescriptor().equals(desc)) {
CodeArea area = editor.getCodeArea();
area.moveTo(child.range().start());
area.showParagraphAtCenter(area.getCurrentParagraph());
return;
}
}
}
}
}
}

@Override
public void onUpdatePath(@Nonnull PathNode<?> path) {
// When we initially set a path, if it's for a class:
// - install the fields/methods navigation side-tab
// - setup handling for what we currently have selected
if (this.path == null && path instanceof ClassPathNode classPathNode) {
// Show declared fields/methods
FieldsAndMethodsPane fieldsAndMethodsPane = fieldsAndMethodsPaneProvider.get();
fieldsAndMethodsPane.setupSelectionNavigationListener(this);
addSideTab(new BoundTab(Lang.getBinding("fieldsandmethods.title"),
new IconView(Icons.getImage(Icons.FIELD_N_METHOD)),
fieldsAndMethodsPane
));
fieldsAndMethodsPane.onUpdatePath(path);

// When the caret position updates, notify the tool tabs of the newly selected member.
editor.getCaretPosEventStream().addObserver(e -> {
if (lastConcreteAst == null)
return;
ClassInfo declaringClass = classPathNode.getValue();
int caret = editor.getCodeArea().getCaretPosition();
for (ASTElement root : lastConcreteAst) {
if (root instanceof ASTClass astClass) {
for (ASTElement child : astClass.children()) {
ClassMember classMember;
if (child instanceof ASTMethod astMethod && astMethod.range().within(caret)) {
String name = astMethod.getName().literal();
String desc = astMethod.getDescriptor().literal();
classMember = declaringClass.getDeclaredMethod(name, desc);
} else if (child instanceof ASTField astField && astField.range().within(caret)) {
String name = astField.getName().literal();
String desc = astField.getDescriptor().literal();
classMember = declaringClass.getDeclaredField(name, desc);
} else {
continue;
}

if (classMember != null) {
assemblerToolTabs.onUpdatePath(classPathNode.child(classMember));
} else {
assemblerToolTabs.onUpdatePath(path);
}
assemblerToolTabs.consumeClass(lastAssembledClassRepresentation, lastAssembledClass);
return;
}
}
}
});
}

this.path = path;

// Update UI state
if (!updateLock.get()) {
pipeline = pipelineManager.getPipeline(path);

Expand Down
65 changes: 44 additions & 21 deletions ...ui/src/main/java/software/coley/recaf/ui/pane/editing/assembler/JvmStackAnalysisPane.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
import me.darknet.assembler.ast.ASTElement;
import me.darknet.assembler.ast.primitive.ASTEmpty;
import me.darknet.assembler.ast.primitive.ASTInstruction;
import me.darknet.assembler.ast.specific.ASTClass;
import me.darknet.assembler.ast.specific.ASTMethod;
import me.darknet.assembler.compile.analysis.*;
import me.darknet.assembler.compile.analysis.frame.Frame;
Expand Down Expand Up @@ -86,28 +87,28 @@ private void updateTable() {
stackTable.setDisable(false);
varTable.setDisable(false);

// Compute what instruction index the caret is at
// Compute what instruction index the caret is at.
int insnIndex = -1;
int pos = editor.getCodeArea().getCaretPosition();
for (ASTElement element : astElements) {
if (element instanceof ASTMethod method && method.range().within(pos)) {
List<ASTInstruction> instructions = method.code().instructions();
int paragraph = editor.getCodeArea().getCurrentParagraph();
int result = Collections.binarySearch(instructions, new ASTEmpty(new Token(
new Range(pos, pos + 1),
new Location(paragraph, 0, null),
TokenType.IDENTIFIER,
"."
)), (o1, o2) -> {
Location l1 = o1.location();
Location l2 = o2.location();
if (l1 == null) return 1;
else if (l2 == null) return -1;
return Objects.compare(l1, l2, Comparator.naturalOrder());
});
if (result < 0) result = -result;
insnIndex = Math.min(instructions.size() - 1, result + 1);
break;
findIndex:
{
for (ASTElement astElement : astElements) {
if (astElement instanceof ASTMethod astMethod) {
int index = getSelectedInsnIndexOfMethod(astMethod);
if (index >= 0) {
insnIndex = index;
break;
}
} else if (astElement instanceof ASTClass astClass) {
for (ASTElement child : astClass.children()) {
if (child instanceof ASTMethod astMethod) {
int index = getSelectedInsnIndexOfMethod(astMethod);
if (index >= 0) {
insnIndex = index;
break findIndex;
}
}
}
}
}
}

Expand Down Expand Up @@ -168,6 +169,28 @@ private void updateTable() {
stackTable.getItems().setAll(stackItems);
}

private int getSelectedInsnIndexOfMethod(@Nonnull ASTMethod method) {
int pos = editor.getCodeArea().getCaretPosition();
if (!method.range().within(pos))
return -1;
List<ASTInstruction> instructions = method.code().instructions();
int paragraph = editor.getCodeArea().getCurrentParagraph();
int result = Collections.binarySearch(instructions, new ASTEmpty(new Token(
new Range(pos, pos + 1),
new Location(paragraph, 0, null),
TokenType.IDENTIFIER,
"."
)), (o1, o2) -> {
Location l1 = o1.location();
Location l2 = o2.location();
if (l1 == null) return 1;
else if (l2 == null) return -1;
return Objects.compare(l1, l2, Comparator.naturalOrder());
});
if (result < 0) result = -result;
return Math.min(instructions.size() - 1, result + 1);
}

private void clearData() {
stackTable.setDisable(true);
varTable.setDisable(true);
Expand Down
53 changes: 34 additions & 19 deletions recaf-ui/src/main/java/software/coley/recaf/ui/pane/editing/assembler/JvmVariablesPane.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@
import me.darknet.assembler.ast.ASTElement;
import me.darknet.assembler.ast.primitive.ASTIdentifier;
import me.darknet.assembler.ast.primitive.ASTInstruction;
import me.darknet.assembler.ast.specific.ASTClass;
import me.darknet.assembler.ast.specific.ASTMember;
import me.darknet.assembler.ast.specific.ASTMethod;
import me.darknet.assembler.compile.analysis.AnalysisResults;
import me.darknet.assembler.compile.analysis.Local;
Expand All @@ -33,6 +35,7 @@
import java.time.Duration;
import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.stream.Stream;

/**
Expand Down Expand Up @@ -91,6 +94,7 @@ protected void updateItem(VariableUsages usages, boolean empty) {
int caret = area.getCaretPosition();
var nextEntry = elementRanges.higherEntry(caret + 1);
if (nextEntry == null) nextEntry = elementRanges.firstEntry();
if (nextEntry == null) return;
Range value = nextEntry.getValue();
area.selectRange(value.start(), value.end());
area.showParagraphAtCenter(area.getCurrentParagraph());
Expand Down Expand Up @@ -145,27 +149,38 @@ private void updateTable() {
variableUsages.put(name, existing.withNewWrite(element));
};
if (astElements != null) {
Consumer<ASTMethod> methodConsumer = astMethod -> {
if (currentMethod != null && !Objects.equals(currentMethod.getName(), astMethod.getName().literal()))
return;
for (ASTIdentifier parameter : astMethod.parameters()) {
String literalName = parameter.literal();
variableUsages.putIfAbsent(literalName, emptyUsage);
}
for (ASTInstruction instruction : astMethod.code().instructions()) {
String insnName = instruction.identifier().content();
boolean isLoad = insnName.endsWith("load");
if (((isLoad || insnName.endsWith("store")) && insnName.charAt(1) != 'a') || insnName.equals("iinc")) {
List<ASTElement> arguments = instruction.arguments();
if (!arguments.isEmpty()) {
ASTElement arg = arguments.get(0);
String varName = arg instanceof ASTIdentifier identifierArg ?
identifierArg.literal() : arg.content();
if (isLoad) {
readUpdater.accept(varName, instruction);
} else {
writeUpdater.accept(varName, instruction);
}
}
}
}
};
for (ASTElement astElement : astElements) {
if (astElement instanceof ASTMethod astMethod) {
for (ASTIdentifier parameter : astMethod.parameters()) {
String literalName = parameter.literal();
variableUsages.putIfAbsent(literalName, emptyUsage);
}
for (ASTInstruction instruction : astMethod.code().instructions()) {
String insnName = instruction.identifier().content();
boolean isLoad = insnName.endsWith("load");
if (((isLoad || insnName.endsWith("store")) && insnName.charAt(1) != 'a') || insnName.equals("iinc")) {
List<ASTElement> arguments = instruction.arguments();
if (arguments.size() > 0) {
ASTElement arg = arguments.get(0);
String varName = arg instanceof ASTIdentifier identifierArg ?
identifierArg.literal() : arg.content();
if (isLoad) {
readUpdater.accept(varName, instruction);
} else {
writeUpdater.accept(varName, instruction);
}
}
methodConsumer.accept(astMethod);
} else if (astElement instanceof ASTClass astClass) {
for (ASTElement child : astClass.children()) {
if (child instanceof ASTMethod astMethod) {
methodConsumer.accept(astMethod);
}
}
}
Expand Down

0 comments on commit e67f619

Please sign in to comment.