Performance and fidelity improvements

Author: Your Name

lightkeeper/src/main/java/lightkeeper/controller/Controller.java

@@ -1,6 +1,5 @@
 package lightkeeper.controller;
 
-import java.awt.Color;
 import java.io.File;
 import java.io.IOException;
 import java.util.List;

lightkeeper/src/main/java/lightkeeper/controller/DisassemblyController.java

@@ -1,19 +1,17 @@
 package lightkeeper.controller;
 
-import java.awt.Color;
 import java.math.BigInteger;
+import java.util.Collection;
 
 import docking.widgets.fieldpanel.LayoutModel;
 import docking.widgets.fieldpanel.listener.IndexMapper;
 import docking.widgets.fieldpanel.listener.LayoutModelListener;
 import generic.theme.GColor;
-import ghidra.app.decompiler.ClangLine;
 import ghidra.app.decompiler.ClangToken;
 import ghidra.app.decompiler.component.DecompilerPanel;
 import ghidra.app.decompiler.component.DecompilerUtils;
 import ghidra.app.plugin.core.decompile.DecompilerActionContext;
 import ghidra.app.plugin.core.decompile.DecompilerProvider;
-import ghidra.program.model.address.AddressRange;
 import ghidra.util.exception.CancelledException;
 import ghidra.util.task.Task;
 import ghidra.util.task.TaskLauncher;
@@ -59,31 +57,9 @@ public void dataChanged(BigInteger start, BigInteger end) {
 		}
 	}
 
-	public void updateCoverage(TaskMonitor monitor, DecompilerPanel dpanel) throws CancelledException {
-		var api = plugin.getApi();
-		if (api == null) {
-			return;
-		}
-
-		var program = api.getCurrentProgram();
-
-		for (ClangLine line : dpanel.getLines()) {
-			for (ClangToken token : line.getAllTokens()) {
-				monitor.checkCancelled();
-				var address = DecompilerUtils.getClosestAddress(program, token);
-				if (address == null) {
-					continue;
-				}
-				for (AddressRange range : model.getModelData()) {
-					monitor.checkCancelled();
-					if (!range.contains(address)) {
-						continue;
-					}
-
-					token.setHighlight(this.color);
-				}
-			}
-		}
+	public void updateCoverage(TaskMonitor monitor, DecompilerPanel dpanel) {
+		Collection<ClangToken> tokens = DecompilerUtils.getTokens(dpanel.getLayoutController().getRoot(), model.getModelData());
+		tokens.stream().forEach(t -> t.setHighlight(this.color));
 	}
 
 	public void updateCoverageTask(DecompilerPanel dpanel) {

lightkeeper/src/main/java/lightkeeper/model/coverage/CoverageModel.java

@@ -14,6 +14,7 @@
 import ghidra.program.model.address.AddressOverflowException;
 import ghidra.program.model.address.AddressRange;
 import ghidra.program.model.address.AddressRangeImpl;
+import ghidra.program.model.address.AddressSet;
 import ghidra.util.exception.CancelledException;
 import ghidra.util.task.TaskMonitor;
 import lightkeeper.LightKeeperPlugin;
@@ -22,10 +23,10 @@
 import lightkeeper.io.module.ModuleEntry;
 import lightkeeper.model.AbstractCoverageModel;
 
-public class CoverageModel extends AbstractCoverageModel<DynamoRioFile, Set<AddressRange>> {
+public class CoverageModel extends AbstractCoverageModel<DynamoRioFile, AddressSet> {
 	protected List<CoverageListRow> rows = new ArrayList<>();
 	protected Map<DynamoRioFile, HashSet<AddressRange>> map = new HashMap<>();
-	protected Set<AddressRange> ranges = new HashSet<>();
+	protected AddressSet ranges = new AddressSet();
 
 	public CoverageModel(LightKeeperPlugin plugin) {
 		super(plugin);
@@ -35,7 +36,7 @@ public CoverageModel(LightKeeperPlugin plugin) {
 	public void clear(TaskMonitor monitor) throws CancelledException {
 		rows = new ArrayList<>();
 		map = new HashMap<>();
-		ranges = new HashSet<>();
+		ranges = new AddressSet();
 		notifyUpdate(monitor);
 	}
 
@@ -93,7 +94,7 @@ protected Set<Integer> getSelectedModuleIds(List<ModuleEntry> selectedModules) {
 	@Override
 	public void update(TaskMonitor monitor) throws CancelledException, IOException {
 		try {
-			ranges = new HashSet<>();
+			ranges = new AddressSet();
 			for (CoverageListRow row : rows) {
 				var state = row.getState();
 				if (state == CoverageListState.IGNORED) {
@@ -140,7 +141,7 @@ public void update(TaskMonitor monitor) throws CancelledException, IOException {
 					.map(r -> map.get(r.file)).flatMap(HashSet::stream).collect(Collectors.toSet());
 
 			added.removeAll(subtracted);
-			ranges = added;
+			added.stream().forEach(r -> ranges.add(r));
 
 			notifyUpdate(monitor);
 		} catch (AddressOverflowException e) {
@@ -153,7 +154,7 @@ public List<CoverageListRow> getFileData() {
 	}
 
 	@Override
-	public Set<AddressRange> getModelData() {
+	public AddressSet getModelData() {
 		return ranges;
 	}
 }

lightkeeper/src/main/java/lightkeeper/model/instruction/CoverageInstructionModel.java

@@ -1,75 +1,51 @@
 package lightkeeper.model.instruction;
 
-import java.util.HashSet;
-import java.util.Set;
-
-import ghidra.program.model.address.AddressRange;
 import ghidra.program.model.address.AddressRangeImpl;
+import ghidra.program.model.address.AddressSet;
 import ghidra.util.exception.CancelledException;
 import ghidra.util.task.TaskMonitor;
 import lightkeeper.LightKeeperPlugin;
 import lightkeeper.model.AbstractCoverageModel;
 import lightkeeper.model.ICoverageModelListener;
 import lightkeeper.model.coverage.CoverageModel;
 
-public class CoverageInstructionModel extends AbstractCoverageModel<Set<AddressRange>, Set<AddressRange>>
+public class CoverageInstructionModel extends AbstractCoverageModel<AddressSet, AddressSet>
 		implements ICoverageModelListener {
 	protected CoverageModel coverage;
-	protected Set<AddressRange> modelRanges;
-	protected Set<AddressRange> hits = new HashSet<>();
+	protected AddressSet modelRanges;
+	protected AddressSet hits = new AddressSet();
 
 	public CoverageInstructionModel(LightKeeperPlugin plugin, CoverageModel coverage) {
 		super(plugin);
 		this.coverage = coverage;
 	}
 
 	@Override
-	public void load(Set<AddressRange> ranges) {
+	public void load(AddressSet ranges) {
 		modelRanges = ranges;
 	}
 
 	@Override
 	public void update(TaskMonitor monitor) throws CancelledException {
-		hits = new HashSet<>();
+		hits = new AddressSet();
 		if (modelRanges == null) {
 			return;
 		}
 
-		var api = plugin.getApi();
-		var listing = api.getCurrentProgram().getListing();
-
-		for (AddressRange range : modelRanges) {
-			monitor.checkCancelled();
-			var iterator = listing.getInstructions(range.getMinAddress(), true);
-			while (iterator.hasNext()) {
-				var instruction = iterator.next();
-
-				if (instruction.getMaxAddress().compareTo(range.getMaxAddress()) > 0) {
-					break;
-				}
-
-				var instructionStart = instruction.getAddress();
-				long length = instruction.getLength();
-				if (length > 0) {
-					length--;
-				}
-				var instructionEnd = instructionStart.add(length);
-				AddressRange instructionRange = new AddressRangeImpl(instructionStart, instructionEnd);
-				hits.add(instructionRange);
-			}
-		}
+		var instructions = plugin.getApi().getCurrentProgram().getListing().getInstructions(modelRanges, true);
+		instructions.forEach(i -> hits.add(new AddressRangeImpl(i.getMinAddress(), i.getMaxAddress())));
 		notifyUpdate(monitor);
 	}
 
 	@Override
 	public void clear(TaskMonitor monitor) throws CancelledException {
 		modelRanges = null;
-		hits = new HashSet<>();
+		hits = new AddressSet();
 		notifyUpdate(monitor);
 	}
 
 	@Override
-	public Set<AddressRange> getModelData() {
+	public AddressSet getModelData() {
 		return hits;
 	}