From 4bc3afe4ff08b9f0c08952ec7f6e0ac930280cc5 Mon Sep 17 00:00:00 2001
From: asie
Date: Sat, 8 Dec 2018 11:21:18 +0100
Subject: add barebones plugin framework, cleanup

---
 .../java/cuchaz/enigma/analysis/SourceIndex.java   | 36 +++++++++++++++++++---
 1 file changed, 31 insertions(+), 5 deletions(-)

(limited to 'src/main/java/cuchaz/enigma/analysis/SourceIndex.java')

diff --git a/src/main/java/cuchaz/enigma/analysis/SourceIndex.java b/src/main/java/cuchaz/enigma/analysis/SourceIndex.java
index 78195cb..4c84e69 100644
--- a/src/main/java/cuchaz/enigma/analysis/SourceIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/SourceIndex.java
@@ -22,10 +22,7 @@ import com.strobel.decompiler.languages.java.ast.Identifier;
 import com.strobel.decompiler.languages.java.ast.TypeDeclaration;
 import cuchaz.enigma.mapping.entry.Entry;
 
-import java.util.Collection;
-import java.util.List;
-import java.util.Map;
-import java.util.TreeMap;
+import java.util.*;
 import java.util.regex.Pattern;
 
 public class SourceIndex {
@@ -48,9 +45,12 @@ public class SourceIndex {
 		this.tokenToReference = Maps.newTreeMap();
 		this.referenceToTokens = HashMultimap.create();
 		this.declarationToToken = Maps.newHashMap();
-		this.lineOffsets = Lists.newArrayList();
+		calculateLineOffsets();
+	}
 
+	private void calculateLineOffsets() {
 		// count the lines
+		this.lineOffsets = Lists.newArrayList();
 		this.lineOffsets.add(0);
 		for (int i = 0; i < source.length(); i++) {
 			if (source.charAt(i) == '\n') {
@@ -59,6 +59,32 @@ public class SourceIndex {
 		}
 	}
 
+	public void remap(String source, Map<Token, Token> tokenMap) {
+		this.source = source;
+		calculateLineOffsets();
+
+		for (Entry entry : Lists.newArrayList(declarationToToken.keySet())) {
+			Token token = declarationToToken.get(entry);
+			declarationToToken.put(entry, tokenMap.getOrDefault(token, token));
+		}
+
+		for (Token token : Lists.newArrayList(tokenToReference.keySet())) {
+			EntryReference<Entry, Entry> e = tokenToReference.remove(token);
+			tokenToReference.put(tokenMap.getOrDefault(token, token), e);
+		}
+
+		for (EntryReference<Entry, Entry> ref : Lists.newArrayList(referenceToTokens.keySet())) {
+			List<Token> newTokens = new ArrayList<>();
+
+			for (Token token : referenceToTokens.get(ref)) {
+				newTokens.add(tokenMap.getOrDefault(token, token));
+			}
+
+			referenceToTokens.removeAll(ref);
+			referenceToTokens.putAll(ref, newTokens);
+		}
+	}
+
 	public String getSource() {
 		return this.source;
 	}
-- 
cgit v1.2.3