fix line numbers in blocks & refactor test

3 files changed, 92 insertions, 65 deletions

diff --git a/src/main/javassist/compiler/Parser.java b/src/main/javassist/compiler/Parser.java
index 43e8a3ec..7e742a6b 100644
--- a/src/main/javassist/compiler/Parser.java
+++ b/src/main/javassist/compiler/Parser.java
@@ -298,20 +298,22 @@ public final class Parser implements TokenId {
     /* block.statement : "{" statement* "}"
      */
     private Stmnt parseBlock(SymbolTable tbl) throws CompileError {
+        int blockLineNumber = lex.getLineNumber();
         if (lex.get() != '{')
             throw new SyntaxError(lex);
 
         Stmnt body = null;
         SymbolTable tbl2 = new SymbolTable(tbl);
         while (lex.lookAhead() != '}') {
+            int lineNumber = lex.getLineNumber();
             Stmnt s = parseStatement(tbl2);
             if (s != null)
-                body = (Stmnt)ASTList.concat(body, new Stmnt(BLOCK, s, lex.getLineNumber()));
+                body = (Stmnt)ASTList.concat(body, new Stmnt(BLOCK, s, lineNumber));
         }
 
         lex.get();      // '}'
         if (body == null)
-            return new Stmnt(BLOCK, lex.getLineNumber());    // empty block
+            return new Stmnt(BLOCK, blockLineNumber);    // empty block
         return body;
     }
 
@@ -321,7 +323,9 @@ public final class Parser implements TokenId {
     private Stmnt parseIf(SymbolTable tbl) throws CompileError {
         int t = lex.get();      // IF
         ASTree expr = parseParExpression(tbl);
+        int exprLineNum = lex.getLineNumber();
         Stmnt thenp = parseStatement(tbl);
+        int thenLineNum = lex.getLineNumber();
         Stmnt elsep;
         if (lex.lookAhead() == ELSE) {
             lex.get();
@@ -330,7 +334,7 @@ public final class Parser implements TokenId {
         else
             elsep = null;
 
-        return new Stmnt(t, expr, new ASTList(thenp, new ASTList(elsep, lex.getLineNumber()), lex.getLineNumber()), lex.getLineNumber());
+        return new Stmnt(t, expr, new ASTList(thenp, new ASTList(elsep, lex.getLineNumber()), thenLineNum), exprLineNum);
     }
 
     /* while.statement : WHILE "(" expression ")" statement
@@ -340,8 +344,9 @@ public final class Parser implements TokenId {
     {
         int t = lex.get();      // WHILE
         ASTree expr = parseParExpression(tbl);
+        int lineNumber = lex.getLineNumber();
         Stmnt body = parseStatement(tbl);
-        return new Stmnt(t, expr, body, lex.getLineNumber());
+        return new Stmnt(t, expr, body, lineNumber);
     }
 
     /* do.statement : DO statement WHILE "(" expression ")" ";"
@@ -372,6 +377,7 @@ public final class Parser implements TokenId {
         if (lex.get() != '(')
             throw new SyntaxError(lex);
 
+        int expr1LineNumber = lex.getLineNumber();
         if (lex.lookAhead() == ';') {
             lex.get();
             expr1 = null;
@@ -379,6 +385,7 @@ public final class Parser implements TokenId {
         else
             expr1 = parseDeclarationOrExpression(tbl2, true);
 
+        int expr2LineNumber = lex.getLineNumber();
         if (lex.lookAhead() == ';')
             expr2 = null;
         else
@@ -387,6 +394,7 @@ public final class Parser implements TokenId {
         if (lex.get() != ';')
             throw new CompileError("; is missing", lex);
 
+        int expr3LineNumber = lex.getLineNumber();
         if (lex.lookAhead() == ')')
             expr3 = null;
         else
@@ -397,7 +405,7 @@ public final class Parser implements TokenId {
 
         Stmnt body = parseStatement(tbl2);
         return new Stmnt(t, expr1, new ASTList(expr2,
-                                               new ASTList(expr3, body, lex.getLineNumber()), lex.getLineNumber()), lex.getLineNumber());
+                                               new ASTList(expr3, body, expr3LineNumber), expr2LineNumber), expr1LineNumber);
     }
 
     /* switch.statement : SWITCH "(" expression ")" "{" switch.block "}"
diff --git a/src/test/javassist/LineNumberTest.java b/src/test/javassist/LineNumberTest.java
new file mode 100644
index 00000000..158efd01
--- /dev/null
+++ b/src/test/javassist/LineNumberTest.java
@@ -0,0 +1,79 @@
+package javassist;
+
+import junit.framework.TestCase;
+
+public class LineNumberTest extends TestCase {
+    private final ClassPool loader = ClassPool.getDefault();
+    private static int classNumber = 0;
+
+    public void testComma() {
+        doTestCompile(String.join("\n",
+                "public void run() {",
+                "    return",
+                "}"), "line 3: syntax error near \"  return\n}\"");
+    }
+
+    public void testException() {
+        doTestRuntime(String.join("\n",
+                "public void run() {",
+                "    throw new java.lang.RuntimeException();",
+                "}"), 0, 5);
+    }
+
+    public void testIf() {
+        doTestRuntime(String.join("\n",
+                "public void run() {",
+                "    if (throwException()) {",
+                "    }",
+                "}"), 1, 5);
+    }
+
+    public void testWhile() {
+        doTestRuntime(String.join("\n",
+                "public void run() {",
+                "    while (throwException()) {",
+                "    }",
+                "}"), 1, 5);
+    }
+
+    public void testFor() {
+        doTestRuntime(String.join("\n",
+                "public void run() {",
+                "    for (; throwException(); ) {",
+                "        ",
+                "    }",
+                "}"), 1, 5);
+    }
+
+    private void doTestCompile(String src, String msg) {
+        CtClass testClass = loader.makeClass("javassist.LineNumberCompileTest" + classNumber++);
+        try {
+            testClass.addMethod(CtMethod.make(src, testClass));
+        } catch (CannotCompileException e) {
+            assertEquals(msg, e.getCause().getMessage());
+            return;
+        }
+        fail("should not happen");
+    }
+
+    private void doTestRuntime(String src, int stackOffset, int lineNumber) {
+        CtClass testClass = loader.makeClass("javassist.LineNumberRuntimeTest" + classNumber++);
+        String test = String.join("\n",
+                "private boolean throwException() {",
+                "    throw new java.lang.RuntimeException();",
+                "}");
+        try {
+            testClass.addInterface(loader.get("java.lang.Runnable"));
+            testClass.addMethod(CtMethod.make(test, testClass));
+            testClass.addMethod(CtMethod.make(src, testClass));
+            Class cls = testClass.toClass(LineNumberTest.class);
+            var runnable = (Runnable) cls.getConstructor().newInstance();
+            runnable.run();
+        } catch (Exception e) {
+            var lineNum = e.getStackTrace()[stackOffset].getLineNumber();
+            assertEquals("Line number should be right", lineNumber, lineNum);
+            return;
+        }
+        fail("should not happen");
+    }
+}
diff --git a/src/test/javassist/bytecode/BytecodeTest.java b/src/test/javassist/bytecode/BytecodeTest.java
index cc02801d..eac420bc 100644
--- a/src/test/javassist/bytecode/BytecodeTest.java
+++ b/src/test/javassist/bytecode/BytecodeTest.java
@@ -302,66 +302,6 @@ public class BytecodeTest extends TestCase {
         assertEquals(20, pc.line);
     }
 
-    public void testLineNumberCompiler() {
-        CtClass testClass = loader.makeClass("javassist.bytecode.LineNumberCompilerClass");
-        String run = String.join("\n",
-                "public void run() {",
-                "    return",
-                "}");
-        try {
-            testClass.addMethod(CtMethod.make(run, testClass));
-        } catch (CannotCompileException e) {
-            assertEquals("line 3: syntax error near \"  return\n}\"", e.getCause().getMessage());
-            return;
-        }
-        fail("should not happen");
-    }
-
-    public void testLineNumberCompiler2() {
-        CtClass testClass = loader.makeClass("javassist.bytecode.LineNumberCompilerClass2");
-        String ctor = String.join("\n",
-                "public LineNumberCompilerClass2() {",
-                "    super();",
-                "}");
-        String run = String.join("\n",
-                "public void run() {",
-                "    return;",
-                "}");
-        String run2 = String.join("\n",
-                "public void run2() {",
-                "    run()",
-                "}");
-        try {
-            testClass.addConstructor(CtNewConstructor.make(ctor, testClass));
-            testClass.addMethod(CtMethod.make(run, testClass));
-            testClass.addMethod(CtMethod.make(run2, testClass));
-        } catch (CannotCompileException e) {
-            assertEquals("line 9: ; is missing", e.getCause().getMessage());
-            return;
-        }
-        fail("should not happen");
-    }
-
-    public void testLineNumberException() {
-        CtClass testClass = loader.makeClass("javassist.bytecode.LineNumberExceptionClass");
-        String run = String.join("\n",
-                "public void run() {",
-                "    throw new java.lang.RuntimeException();",
-                "}");
-        try {
-            testClass.addInterface(loader.get("java.lang.Runnable"));
-            testClass.addMethod(CtMethod.make(run, testClass));
-            Class cls = testClass.toClass(BytecodeTest.class);
-            var runnable = (Runnable) cls.getConstructor().newInstance();
-            runnable.run();
-        } catch (Exception e) {
-            var lineNum = e.getStackTrace()[0].getLineNumber();
-            assertEquals("Line number should be right", 2, lineNum);
-            return;
-        }
-        fail("should not happen");
-    }
-
     public void testRenameClass() throws Exception {
         CtClass cc = loader.get("test1.RenameClass");
         cc.replaceClassName("test1.RenameClass2", "java.lang.String");