core: fix source line number parsing and saving

jadx-core/src/main/java/jadx/core/codegen/CodeWriter.java

@@ -10,6 +10,7 @@ import java.util.Collections;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.Map;
+import java.util.TreeMap;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -172,11 +173,11 @@ public class CodeWriter {
 		return attachAnnotation(obj, new CodePosition(line, offset + 1));
 	}
 
-	private void attachSourceLine(int decompiledLine, int sourceLine) {
-		if (lineMap.isEmpty()) {
-			lineMap = new HashMap<Integer, Integer>();
+	private Object attachAnnotation(Object obj, CodePosition pos) {
+		if (annotations.isEmpty()) {
+			annotations = new HashMap<CodePosition, Object>();
 		}
-		lineMap.put(decompiledLine, sourceLine);
+		return annotations.put(pos, obj);
 	}
 
 	public Map<CodePosition, Object> getAnnotations() {
@@ -184,18 +185,21 @@ public class CodeWriter {
 	}
 
 	public void attachSourceLine(int sourceLine) {
+		if (sourceLine == 0) {
+			return;
+		}
 		attachSourceLine(line, sourceLine);
 	}
 
-	public Map<Integer, Integer> getLineMapping() {
-		return lineMap;
+	private void attachSourceLine(int decompiledLine, int sourceLine) {
+		if (lineMap.isEmpty()) {
+			lineMap = new TreeMap<Integer, Integer>();
 		}
-
-	private Object attachAnnotation(Object obj, CodePosition pos) {
-		if (annotations.isEmpty()) {
-			annotations = new HashMap<CodePosition, Object>();
+		lineMap.put(decompiledLine, sourceLine);
 	}
-		return annotations.put(pos, obj);
+
+	public Map<Integer, Integer> getLineMapping() {
+		return lineMap;
 	}
 
 	public void finish() {

jadx-core/src/main/java/jadx/core/codegen/MethodGen.java

@@ -86,6 +86,7 @@ public class MethodGen {
 			ai = ai.remove(AccessFlags.ACC_PUBLIC);
 		}
 		code.startLine(ai.makeString());
+		code.attachSourceLine(mth.getSourceLine());
 
 		if (classGen.addGenericMap(code, mth.getGenericMap())) {
 			code.add(' ');

jadx-core/src/main/java/jadx/core/codegen/RegionGen.java

@@ -108,6 +108,7 @@ public class RegionGen extends InsnGen {
 		if (newLine) {
 			code.startLine();
 		}
+		code.attachSourceLine(region.getSourceLine());
 		code.add("if (");
 		new ConditionGen(this).add(code, region.getCondition());
 		code.add(") {");

jadx-core/src/main/java/jadx/core/dex/instructions/mods/ConstructorInsn.java

@@ -49,6 +49,7 @@ public class ConstructorInsn extends InsnNode {
 			addArg(invoke.getArg(i));
 		}
 		offset = invoke.getOffset();
+		setSourceLine(invoke.getSourceLine());
 	}
 
 	public MethodInfo getCallMth() {

jadx-core/src/main/java/jadx/core/dex/nodes/parser/DebugInfoParser.java

@@ -75,6 +75,7 @@ public class DebugInfoParser {
 
 		// process '0' instruction
 		addrChange(-1, 1, line);
+		setLine(addr, line);
 
 		int c = section.readByte() & 0xFF;
 		while (c != DBG_END_SEQUENCE) {
@@ -82,6 +83,7 @@ public class DebugInfoParser {
 				case DBG_ADVANCE_PC: {
 					int addrInc = section.readUleb128();
 					addr = addrChange(addr, addrInc, line);
+					setLine(addr, line);
 					break;
 				}
 				case DBG_ADVANCE_LINE: {
@@ -143,9 +145,10 @@ public class DebugInfoParser {
 				default: {
 					if (c >= DBG_FIRST_SPECIAL) {
 						int adjustedOpcode = c - DBG_FIRST_SPECIAL;
-						line += DBG_LINE_BASE + (adjustedOpcode % DBG_LINE_RANGE);
 						int addrInc = adjustedOpcode / DBG_LINE_RANGE;
 						addr = addrChange(addr, addrInc, line);
+						line += DBG_LINE_BASE + (adjustedOpcode % DBG_LINE_RANGE);
+						setLine(addr, line);
 					} else {
 						throw new DecodeException("Unknown debug insn code: " + c);
 					}
@@ -161,6 +164,7 @@ public class DebugInfoParser {
 				setVar(var);
 			}
 		}
+		setSourceLines(addr, insnByOffset.length, line);
 	}
 
 	private int addrChange(int addr, int addrInc, int line) {
@@ -170,7 +174,6 @@ public class DebugInfoParser {
 			if (insn == null) {
 				continue;
 			}
-			insn.setSourceLine(line);
 			for (InsnArg arg : insn.getArguments()) {
 				if (arg.isRegister()) {
 					activeRegisters[((RegisterArg) arg).getRegNum()] = arg;
@@ -181,9 +184,23 @@ public class DebugInfoParser {
 				activeRegisters[res.getRegNum()] = res;
 			}
 		}
+		setSourceLines(addr, newAddr, line);
 		return newAddr;
 	}
 
+	private void setSourceLines(int start, int end, int line) {
+		for (int offset = start + 1; offset < end; offset++) {
+			setLine(offset, line);
+		}
+	}
+
+	private void setLine(int offset, int line) {
+		InsnNode insn = insnByOffset[offset];
+		if (insn != null) {
+			insn.setSourceLine(line);
+		}
+	}
+
 	private void startVar(LocalVar var, int addr, int line) {
 		int regNum = var.getRegNum();
 		LocalVar prev = locals[regNum];

jadx-core/src/main/java/jadx/core/dex/regions/IfRegion.java

@@ -78,6 +78,13 @@ public final class IfRegion extends AbstractRegion {
 		elseRegion = tmp;
 	}
 
+	public int getSourceLine() {
+		if (header.getInstructions().isEmpty()) {
+			return 0;
+		}
+		return header.getInstructions().get(0).getSourceLine();
+	}
+
 	@Override
 	public List<IContainer> getSubBlocks() {
 		if (ternRegion != null) {

jadx-core/src/main/java/jadx/core/dex/visitors/BlockMakerVisitor.java

@@ -488,6 +488,7 @@ public class BlockMakerVisitor extends AbstractVisitor {
 		}
 		insn.copyAttributesFrom(returnInsn);
 		insn.setOffset(returnInsn.getOffset());
+		insn.setSourceLine(returnInsn.getSourceLine());
 		return insn;
 	}
 

jadx-core/src/main/java/jadx/core/dex/visitors/regions/TernaryMod.java

@@ -48,6 +48,7 @@ public class TernaryMod {
 			RegisterArg resArg = t.getResult().getSVar().getUsedInPhi().getResult();
 			TernaryInsn ternInsn = new TernaryInsn(ifRegion.getCondition(),
 					resArg, InsnArg.wrapArg(t), InsnArg.wrapArg(e));
+			ternInsn.setSourceLine(t.getSourceLine());
 			TernaryRegion tern = new TernaryRegion(ifRegion, header);
 			// TODO: add api for replace regions
 			ifRegion.setTernRegion(tern);
@@ -69,6 +70,7 @@ public class TernaryMod {
 			eb.remove(AFlag.RETURN);
 
 			TernaryInsn ternInsn = new TernaryInsn(ifRegion.getCondition(), null, t.getArg(0), e.getArg(0));
+			ternInsn.setSourceLine(t.getSourceLine());
 			InsnNode retInsn = new InsnNode(InsnType.RETURN, 1);
 			retInsn.addArg(InsnArg.wrapArg(ternInsn));
 

jadx-core/src/test/java/jadx/tests/internal/debuginfo/TestLineNumbers2.java

+package jadx.tests.internal.debuginfo;
+
+import jadx.api.InternalJadxTest;
+import jadx.core.codegen.CodeWriter;
+import jadx.core.dex.nodes.ClassNode;
+
+import java.lang.ref.WeakReference;
+import java.util.Map;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+
+public class TestLineNumbers2 extends InternalJadxTest {
+
+	public static class TestCls {
+		private WeakReference<TestCls> f;
+
+		public TestCls(TestCls s) {
+		}
+
+		TestCls test(TestCls s) {
+			TestCls store = f != null ? f.get() : null;
+			if (store == null) {
+				store = new TestCls(s);
+				f = new WeakReference<TestCls>(store);
+			}
+			return store;
+		}
+
+		public Object test2() {
+			return new Object();
+		}
+	}
+
+	@Test
+	public void test() {
+		ClassNode cls = getClassNode(TestCls.class);
+		CodeWriter codeWriter = cls.getCode();
+		String code = codeWriter.toString();
+		System.out.println(code);
+
+		Map<Integer, Integer> lineMapping = codeWriter.getLineMapping();
+		assertEquals("{8=18, 11=22, 13=23, 14=24, 15=28, 17=25, 18=26, 19=28, 22=31, 23=32}",
+				lineMapping.toString());
+	}
+}